add main files

Dockerfile

@@ -0,0 +1,20 @@
+# Base image
+FROM python:3.9-slim-buster
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the requirements.txt file
+COPY requirements.txt .
+
+# Install the dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the code into the container
+COPY . .
+
+# Expose the port
+EXPOSE 7860
+
+# Run the application
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,176 @@
+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 PIL import Image
+
+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
+)
+
+pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    safety_checker=None,
+    torch_dtype=torch.float16,
+).to("cuda")
+pipe.enable_xformers_memory_efficient_attention()
+
+
+def resize_for_condition_image(input_image: Image.Image, resolution: int):
+    input_image = input_image.convert("RGB")
+    W, H = input_image.size
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(round(H / 64.0)) * 64
+    W = int(round(W / 64.0)) * 64
+    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'
+
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860)

requirements.txt

@@ -0,0 +1,9 @@
+torch==1.9.0
+gradio==2.0.8
+Pillow==8.3.2
+qrcode==7.3
+pathlib==1.0.1
+requests==2.26.0
+diffusers==0.0.1
+Flask==2.0.2
+uvicorn