Update app.py

app.py

@@ -13,6 +13,8 @@ import numpy as np
 import cv2
 from pyzxing import BarCodeReader
 import tempfile
+from io import BytesIO
+import base64
 
 from diffusers import (
     StableDiffusionPipeline,
@@ -135,6 +137,32 @@ def scan_qr_code(image, sensitivity=0.5):
     
     return successful_scans
 
+def generate_qr_code(
+    data,
+    version=None,
+    error_correction=qrcode.constants.ERROR_CORRECT_H,
+    box_size=10,
+    border=4,
+    fill_color="black",
+    back_color="white"
+):
+    qr = qrcode.QRCode(
+        version=version,
+        error_correction=error_correction,
+        box_size=box_size,
+        border=border,
+    )
+    qr.add_data(data)
+    qr.make(fit=True)
+    
+    img = qr.make_image(fill_color=fill_color, back_color=back_color)
+    return img
+
+def image_to_base64(img):
+    buffered = BytesIO()
+    img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
 @spaces.GPU()
 def inference(
     qr_code_content: str,
@@ -149,6 +177,12 @@ def inference(
     use_qr_code_as_init_image = True,
     sampler = "DPM++ Karras SDE",
     verify_sensitivity: float = 0.5,
+    qr_version: int = 1,
+    qr_error_correction: str = "H",
+    qr_box_size: int = 10,
+    qr_border: int = 4,
+    qr_fill_color: str = "black",
+    qr_back_color: str = "white",
 ):
     try:
         if prompt is None or prompt == "":
@@ -163,33 +197,39 @@ def inference(
 
         if qr_code_content != "" or qrcode_image.size == (1, 1):
             print("Generating QR Code from content")
-            qr = qrcode.QRCode(
-                version=1,
-                error_correction=qrcode.constants.ERROR_CORRECT_H,
-                box_size=10,
-                border=4,
+            error_correction_map = {
+                "L": qrcode.constants.ERROR_CORRECT_L,
+                "M": qrcode.constants.ERROR_CORRECT_M,
+                "Q": qrcode.constants.ERROR_CORRECT_Q,
+                "H": qrcode.constants.ERROR_CORRECT_H,
+            }
+            qrcode_image = generate_qr_code(
+                qr_code_content,
+                version=qr_version,
+                error_correction=error_correction_map[qr_error_correction],
+                box_size=qr_box_size,
+                border=qr_border,
+                fill_color=qr_fill_color,
+                back_color=qr_back_color,
             )
-            qr.add_data(qr_code_content)
-            qr.make(fit=True)
-
-            qrcode_image = qr.make_image(fill_color="black", back_color="white")
-            qrcode_image = resize_for_condition_image(qrcode_image, 768)
-        else:
-            print("Using QR Code Image")
             qrcode_image = resize_for_condition_image(qrcode_image, 768)
 
-        # hack due to gradio examples
-        init_image = qrcode_image
+        # Generate base64 string of the initial QR code
+        initial_qr_base64 = image_to_base64(qrcode_image)
+
+        # Use QR code as init image if specified
+        if use_qr_code_as_init_image:
+            init_image = qrcode_image
 
         out = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            image=qrcode_image,
-            control_image=qrcode_image,  # type: ignore
-            width=768,  # type: ignore
-            height=768,  # type: ignore
+            image=init_image,
+            control_image=qrcode_image,
+            width=768,
+            height=768,
             guidance_scale=float(guidance_scale),
-            controlnet_conditioning_scale=float(controlnet_conditioning_scale),  # type: ignore
+            controlnet_conditioning_scale=float(controlnet_conditioning_scale),
             generator=generator,
             strength=float(strength),
             num_inference_steps=50,
@@ -198,11 +238,11 @@ def inference(
         # Verify the generated QR code
         verify_result = scan_qr_code(out.images[0], sensitivity=verify_sensitivity)
         
-        return out.images[0], verify_result
+        return out.images[0], verify_result, initial_qr_base64
     except Exception as e:
         print(f"Error in inference: {str(e)}")
         # Return a blank image in case of an error
-        return Image.new('RGB', (768, 768), color='white'), []
+        return Image.new('RGB', (768, 768), color='white'), [], ""
 
 
 
@@ -374,7 +414,6 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as blocks:
                     maximum=1.0,
                     step=0.05,
                     value=0.5,
-                    visible=True,
                     label="Verify Sensitivity",
                 )
                 gr.Markdown(
@@ -389,6 +428,8 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as blocks:
                 
         with gr.Column():
             result_image = gr.Image(label="Your Artistic QR Code")
+            initial_qr_image = gr.Image(label="Initial QR Code", visible=False)
+            download_initial_qr = gr.Button("Download Initial QR Code")
             scan_button = gr.Button("Scan QR Code")
             scan_result = gr.Textbox(label="Scan Result", interactive=False)
             
@@ -405,6 +446,14 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as blocks:
                 """
             )
 
+        with gr.Accordion("QR Code Customization", open=False):
+            qr_version = gr.Slider(minimum=1, maximum=40, step=1, value=1, label="QR Code Version")
+            qr_error_correction = gr.Dropdown(choices=["L", "M", "Q", "H"], value="H", label="Error Correction Level")
+            qr_box_size = gr.Slider(minimum=1, maximum=50, step=1, value=10, label="Box Size")
+            qr_border = gr.Slider(minimum=0, maximum=10, step=1, value=4, label="Border Size")
+            qr_fill_color = gr.ColorPicker(label="QR Code Color", value="#000000")
+            qr_back_color = gr.ColorPicker(label="Background Color", value="#FFFFFF")
+
     def scan_and_display(image):
         if image is None:
             return "No image to scan"
@@ -424,7 +473,9 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as blocks:
         outputs=[scan_result]
     )
 
-            
+    def update_initial_qr(base64_string):
+        return gr.Image.update(value=base64_string, visible=True)
+
     run_btn.click(
         inference,
         inputs=[
@@ -440,10 +491,19 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as blocks:
             use_qr_code_as_init_image,
             sampler,
             verify_sensitivity,
+            qr_version,
+            qr_error_correction,
+            qr_box_size,
+            qr_border,
+            qr_fill_color,
+            qr_back_color,
         ],
-        outputs=[result_image, scan_result],
+        outputs=[result_image, scan_result, initial_qr_image],
         concurrency_limit=20
     )
 
+    download_initial_qr.click(lambda x: x, inputs=[initial_qr_image], outputs=[initial_qr_image])
+
+
 blocks.queue(max_size=20,api_open=False)
 blocks.launch(share=bool(os.environ.get("SHARE", False)), show_api=True)