Update app.py

app.py

@@ -26,6 +26,8 @@ from diffusers import (
     EulerDiscreteScheduler,
 )
 
+
+
 qrcode_generator = qrcode.QRCode(
     version=1,
     error_correction=qrcode.ERROR_CORRECT_H,
@@ -33,25 +35,38 @@ qrcode_generator = qrcode.QRCode(
     border=4,
 )
 
-controlnet = ControlNetModel.from_pretrained(
-    #"monster-labs/control_v1p_sd15_qrcode_monster",
-    #"DionTimmer/controlnet_qrcode-control_v1p_sd15",
-    "DionTimmer/controlnet_qrcode",
-    torch_dtype=torch.float16
-).to("cuda")
-
-pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-    #"digiplay/GhostMixV1.2VAE",
-    #"benjamin-paine/stable-diffusion-v1-5",
-    "ckpt/sd15",
-    controlnet = controlnet,
-    torch_dtype = torch.float16,
-    safety_checker = None,
-).to("cuda")
-#pipe.enable_xformers_memory_efficient_attention()
-#pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-#pipe.enable_model_cpu_offload()
 
+# Define available models
+CONTROLNET_MODELS = {
+    "QR Code Monster": "monster-labs/control_v1p_sd15_qrcode_monster",
+    "QR Code": "DionTimmer/controlnet_qrcode",
+    # Add more ControlNet models here
+}
+
+DIFFUSION_MODELS = {
+    "Stable Diffusion v1.5": "runwayml/stable-diffusion-v1-5",
+    "GhostMix v1.2": "digiplay/GhostMixV1.2VAE",
+    # Add more diffusion models here
+}
+
+def load_models(controlnet_model, diffusion_model):
+    hf_token = get_hf_token()
+    
+    controlnet = ControlNetModel.from_pretrained(
+        CONTROLNET_MODELS[controlnet_model],
+        use_auth_token=hf_token,
+        torch_dtype=torch.float16
+    ).to("cuda")
+
+    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+        DIFFUSION_MODELS[diffusion_model],
+        controlnet=controlnet,
+        use_auth_token=hf_token,
+        torch_dtype=torch.float16,
+        safety_checker=None,
+    ).to("cuda")
+    
+    return pipe
 
 def resize_for_condition_image(input_image: Image.Image, resolution: int):
     input_image = input_image.convert("RGB")
@@ -148,8 +163,14 @@ def inference(
     bg_color: str = "white",
     qr_color: str = "black",
     invert_final_image: bool = False,
+    invert_init_image: bool = False,  # New parameter
+    controlnet_model: str = "QR Code Monster",
+    diffusion_model: str = "Stable Diffusion v1.5",
 ):
     try:
+        # Load models based on user selection
+        pipe = load_models(controlnet_model, diffusion_model)
+        
         if prompt is None or prompt == "":
             raise gr.Error("Prompt is required")
 
@@ -158,7 +179,9 @@ def inference(
 
         pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)
 
-        generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()
+        if seed == -1:
+            seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        generator = torch.manual_seed(seed)
 
         if qr_code_content != "" or qrcode_image.size == (1, 1):
             print("Generating QR Code from content")
@@ -191,6 +214,10 @@ def inference(
         if init_image is not None:
             strength = min(strength, 0.8)  # Cap strength at 0.8 when using init_image
 
+        # Invert init_image if requested
+        if invert_init_image and init_image is not None:
+            init_image = invert_image(init_image)
+
         out = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -210,13 +237,19 @@ def inference(
         if invert_final_image:
             final_image = invert_image(final_image)
         
-        return final_image
+        return final_image, seed
     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 a blank image and -1 as seed in case of an error
+        return Image.new('RGB', (768, 768), color='white'), -1
 
+def invert_init_image_display(image):
+    if image is None:
+        return None
+    inverted = invert_image(image)
+    if isinstance(inverted, np.ndarray):
+        return Image.fromarray(inverted)
+    return inverted
 
 with gr.Blocks(theme='Hev832/Applio') as blocks:
     gr.Markdown(
@@ -261,7 +294,7 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
             prompt = gr.Textbox(
                 label="Artistic Prompt",
                 placeholder="Describe the style or theme for your QR code art",
-                info="Be specific and creative! This guides the AI in creating your unique QR art.",
+                info="Be specific and creative! This guides the AI in creating your unique QR code art.",
             )
             negative_prompt = gr.Textbox(
                 label="Elements to Avoid",
@@ -277,7 +310,13 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
             )
             with gr.Accordion(label="Init Images (Optional)", open=False, visible=True) as init_image_acc:
                 init_image = gr.Image(label="Init Image (Optional). Leave blank to generate image with AI", type="pil")
-
+                with gr.Row():
+                    invert_init_image = gr.Checkbox(
+                        label="Invert Init Image", 
+                        value=False,
+                        info="Check this to invert the colors of the init image"
+                    )
+                    invert_init_image_button = gr.Button("Invert Init Image")
 
             with gr.Accordion("Advanced Art Controls", open=True):
             
@@ -396,11 +435,25 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     value=False,
                     info="Check this to invert the colors of the final image"
                 )
+            with gr.Accordion("Model Selection", open=True):
+                controlnet_model_dropdown = gr.Dropdown(
+                    choices=list(CONTROLNET_MODELS.keys()),
+                    value="QR Code Monster",
+                    label="ControlNet Model",
+                    info="Select the ControlNet model for QR code generation"
+                )
+                diffusion_model_dropdown = gr.Dropdown(
+                    choices=list(DIFFUSION_MODELS.keys()),
+                    value="Stable Diffusion v1.5",
+                    label="Diffusion Model",
+                    info="Select the main diffusion model for image generation"
+                )
             with gr.Row():
                 run_btn = gr.Button("🎨 Create Your QR Art", variant="primary")
                 
         with gr.Column():
             result_image = gr.Image(label="Your Artistic QR Code")
+            used_seed = gr.Number(label="Seed Used", interactive=False)
             with gr.Row():
                 scan_button = gr.Button("Scan QR Code")
                 invert_button = gr.Button("Invert Image")
@@ -446,6 +499,12 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
         outputs=[result_image]
     )
 
+    invert_init_image_button.click(
+        invert_init_image_display,
+        inputs=[init_image],
+        outputs=[init_image]
+    )
+
     run_btn.click(
         inference,
         inputs=[
@@ -463,8 +522,11 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
             bg_color,
             qr_color,
             invert_final_image,
+            invert_init_image,
+            controlnet_model_dropdown,
+            diffusion_model_dropdown,
         ],
-        outputs=[result_image],
+        outputs=[result_image, used_seed],
         concurrency_limit=20
     )