Update app.py

app.py

@@ -1,53 +1,26 @@
-import torch
-import gradio as gr
-from PIL import Image
-import qrcode
-from pathlib import Path
-from multiprocessing import cpu_count
-import requests
-import io
-import os
-
-
+from flask import Flask, request, jsonify
 from PIL import Image
+import torch
 from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
 from diffusers.utils import load_image
 
+app = Flask(__name__)
 
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionControlNetImg2ImgPipeline,
-    ControlNetModel,
-    DDIMScheduler,
-    DPMSolverMultistepScheduler,
-    DEISMultistepScheduler,
-    HeunDiscreteScheduler,
-    EulerDiscreteScheduler,
-)
-
-qrcode_generator = qrcode.QRCode(
-    version=1,
-    error_correction=qrcode.ERROR_CORRECT_H,
-    box_size=10,
-    border=4,
-)
-
-controlnet = ControlNetModel.from_pretrained(
-    "DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch.float16
-)
+controlnet = ControlNetModel.from_pretrained("DionTimmer/controlnet_qrcode-control_v1p_sd15",
+                                             torch_dtype=torch.float16)
 
 pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     controlnet=controlnet,
     safety_checker=None,
-    torch_dtype=torch.float16,
-)  #.to("cuda")
+    torch_dtype=torch.float16
+)
+
 pipe.enable_xformers_memory_efficient_attention()
 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 pipe.enable_model_cpu_offload()
 
-
-def resize_for_condition_image(input_image: Image.Image, resolution: int):
+def resize_for_condition_image(input_image: Image, resolution: int):
     input_image = input_image.convert("RGB")
     W, H = input_image.size
     k = float(resolution) / min(H, W)
@@ -58,126 +31,38 @@ def resize_for_condition_image(input_image: Image.Image, resolution: int):
     img = input_image.resize((W, H), resample=Image.LANCZOS)
     return img
 
-
-SAMPLER_MAP = {
-    "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
-    "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),
-    "Heun": lambda config: HeunDiscreteScheduler.from_config(config),
-    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
-    "DDIM": lambda config: DDIMScheduler.from_config(config),
-    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),
-}
-
-
-def inference(
-    qr_code_content: str,
-    prompt: str,
-    negative_prompt: str,
-    guidance_scale: float = 10.0,
-    controlnet_conditioning_scale: float = 2.0,
-    strength: float = 0.8,
-    seed: int = -1,
-    init_image: Image.Image | None = None,
-    qrcode_image: Image.Image | None = None,
-    use_qr_code_as_init_image = True,
-    sampler = "DPM++ Karras SDE",
-):
-    if prompt is None or prompt == "":
-        raise gr.Error("Prompt is required")
-
-    if qrcode_image is None and qr_code_content == "":
-        raise gr.Error("QR Code Image or QR Code Content is required")
-
-    pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)
-
-    generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()
-
-    if qr_code_content != "" or qrcode_image.size == (1, 1):
-        qr = qrcode.QRCode(
-            version=1,
-            error_correction=qrcode.constants.ERROR_CORRECT_H,
-            box_size=10,
-            border=4,
-        )
-        qr.add_data(qr_code_content)
-        qr.make(fit=True)
-        qrcode_image = qr.make_image(fill_color="black", back_color="white")
-
-    if init_image is None:
-        if use_qr_code_as_init_image:
-            init_image = qrcode_image.convert("RGB")
-
-    resolution = controlnet.config.resolution
-    qrcode_image = resize_for_condition_image(qrcode_image, resolution)
-    if init_image is not None:
-        init_image = init_image.convert("RGB")
-        init_image = resize_for_condition_image(init_image, resolution)
-        init_image = torch.nn.functional.interpolate(
-            torch.nn.functional.to_tensor(init_image).unsqueeze(0),
-            size=(resolution, resolution),
-            mode="bilinear",
-            align_corners=False,
-        )[0].unsqueeze(0)
-    else:
-        init_image = torch.zeros(
-            (1, 3, resolution, resolution), device=pipe.device
-        ).to(dtype=torch.float32)
-
-    with torch.no_grad():
-        result_image = pipe(
-            qr_code_condition=qrcode_image,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            init_image=init_image,
-            strength=strength,
-            guidance_scale=guidance_scale,
-            controlnet_conditioning_scale=controlnet_conditioning_scale,
-            disable_progress_bar=True,
-            seed=generator,
-        ).cpu()
-
-    result_image = (
-        result_image.clamp(-1, 1).squeeze().permute(1, 2, 0).numpy() * 255
-    )
-    result_image = Image.fromarray(result_image.astype("uint8"))
-
-    return result_image
-
-
-app = Flask(__name__)
-
-@app.route('/generate_qr_code', methods=['POST'])
-def generate_qr_code():
-    qr_code_content = request.json['qr_code_content']
-    prompt = request.json['prompt']
-    negative_prompt = request.json['negative_prompt']
-    guidance_scale = float(request.json.get('guidance_scale', 10.0))
-    controlnet_conditioning_scale = float(request.json.get('controlnet_conditioning_scale', 2.0))
-    strength = float(request.json.get('strength', 0.8))
-    seed = int(request.json.get('seed', -1))
-    init_image = None
-    qrcode_image = None
-    use_qr_code_as_init_image = request.json.get('use_qr_code_as_init_image', True)
-    sampler = request.json.get('sampler', 'DPM++ Karras SDE')
-
-    try:
-        result_image = inference(qr_code_content, prompt, negative_prompt, guidance_scale,
-                                 controlnet_conditioning_scale, strength, seed, init_image,
-                                 qrcode_image, use_qr_code_as_init_image, sampler)
-
-        image_bytes = io.BytesIO()
-        result_image.save(image_bytes, format='PNG')
-        image_base64 = base64.b64encode(image_bytes.getvalue()).decode('utf-8')
-
-        return jsonify({'image_base64': image_base64})
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-
-@app.route('/health', methods=['GET'])
-def health_check():
-    return 'OK'
-
+@app.route('/generate_image', methods=['POST'])
+def generate_image():
+    # Get input parameters from the request
+    prompt = request.json.get('prompt')
+    negative_prompt = request.json.get('negative_prompt')
+    image_url = request.json.get('image_url')
+    control_image_url = request.json.get('control_image_url')
+
+    # Load the images from URLs
+    source_image = load_image(image_url)
+    init_image = load_image(control_image_url)
+    
+    # Resize images for conditioning
+    condition_image = resize_for_condition_image(source_image, 768)
+    init_image = resize_for_condition_image(init_image, 768)
+    
+    # Generate the image using the pipeline
+    generator = torch.manual_seed(123121231)
+    image = pipe(prompt=prompt,
+                 negative_prompt=negative_prompt, 
+                 image=init_image,
+                 control_image=condition_image,
+                 width=768,
+                 height=768,
+                 guidance_scale=20,
+                 controlnet_conditioning_scale=1.5,
+                 generator=generator,
+                 strength=0.9, 
+                 num_inference_steps=150)
+    
+    # Return the generated image
+    return jsonify({'image': image.images[0]})
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    app.run()