Update app.py

app.py

@@ -12,6 +12,7 @@ import spaces
 import numpy as np
 import cv2
 from pyzxing import BarCodeReader
+from PIL import ImageOps
 
 from diffusers import (
     StableDiffusionPipeline,
@@ -109,6 +110,8 @@ def scan_qr_code(image):
     
     return None
 
+def invert_image(image):
+    return ImageOps.invert(image)
 
 @spaces.GPU()
 def inference(
@@ -123,6 +126,9 @@ def inference(
     qrcode_image: Image.Image | None = None,
     use_qr_code_as_init_image = True,
     sampler = "DPM++ Karras SDE",
+    bg_color: str = "white",
+    qr_color: str = "black",
+    invert_final_image: bool = False,
 ):
     try:
         if prompt is None or prompt == "":
@@ -146,7 +152,7 @@ def inference(
             qr.add_data(qr_code_content)
             qr.make(fit=True)
 
-            qrcode_image = qr.make_image(fill_color="black", back_color="white")
+            qrcode_image = qr.make_image(fill_color=qr_color, back_color=bg_color)
             qrcode_image = resize_for_condition_image(qrcode_image, 768)
         else:
             print("Using QR Code Image")
@@ -179,7 +185,13 @@ def inference(
             strength=float(strength),
             num_inference_steps=50,
         )
-        return out.images[0]
+        
+        final_image = out.images[0]
+        
+        if invert_final_image:
+            final_image = invert_image(final_image)
+        
+        return final_image
     except Exception as e:
         print(f"Error in inference: {str(e)}")
         # Return a blank image in case of an error
@@ -349,6 +361,22 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     Try **-1** if you want to explore and generate different designs. If you find something you really love, write down the seed number and use it again to recreate the same design.
                     """
                 )
+            with gr.Accordion("QR Code Customization", open=True):
+                bg_color = gr.ColorPicker(
+                    label="Background Color", 
+                    value="#FFFFFF",
+                    info="Choose the background color for the QR code"
+                )
+                qr_color = gr.ColorPicker(
+                    label="QR Code Color", 
+                    value="#000000",
+                    info="Choose the color for the QR code pattern"
+                )
+                invert_final_image = gr.Checkbox(
+                    label="Invert Final Image", 
+                    value=False,
+                    info="Check this to invert the colors of the final image"
+                )
             with gr.Row():
                 run_btn = gr.Button("🎨 Create Your QR Art", variant="primary")
                 
@@ -401,6 +429,9 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
             qr_code_image,
             use_qr_code_as_init_image,
             sampler,
+            bg_color,
+            qr_color,
+            invert_final_image,
         ],
         outputs=[result_image],
         concurrency_limit=20