Update app.py

app.py

@@ -36,7 +36,6 @@ load_dotenv()
 USERNAME = os.getenv("USERNAME")
 PASSWORD = os.getenv("PASSWORD")
 
-
 qrcode_generator = qrcode.QRCode(
     version=1,
     error_correction=qrcode.ERROR_CORRECT_H,
@@ -75,7 +74,7 @@ def load_models_on_launch():
     loaded_controlnet = ControlNetModel.from_pretrained(
         controlnet_path,
         torch_dtype=torch.float16
-    ).to("cuda")
+    ).to("mps")
 
     diffusion_path = snapshot_download(DIFFUSION_MODELS["GhostMix"])
     loaded_pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
@@ -83,7 +82,7 @@ def load_models_on_launch():
         controlnet=loaded_controlnet,
         torch_dtype=torch.float16,
         safety_checker=None,
-    ).to("cuda")
+    ).to("mps")
     print("Models loaded successfully!")
 
 # Modify the load_models function to use global variables
@@ -200,9 +199,9 @@ def inference(
     qr_code_content: str,
     prompt: str,
     negative_prompt: str,
-    guidance_scale: float = 15.0,
-    controlnet_conditioning_scale: float = 1.5,
-    strength: float = 0.6,
+    guidance_scale: float = 15.0,  # Increased from 10.0 to 15.0
+    controlnet_conditioning_scale: float = 1.5,  # Adjusted from 2.0 to 1.5
+    strength: float = 0.6,  # Reduced from 0.8 to 0.6
     seed: int = -1,
     init_image: Image.Image | None = None,
     qrcode_image: Image.Image | None = None,
@@ -215,14 +214,13 @@ def inference(
     controlnet_model: str = "QR Code Monster",
     diffusion_model: str = "GhostMix",
     reference_image_strength: float = 0.6,
-    progress: gr.Progress = gr.Progress()
 ):
     try:
-        progress(0, desc="Initializing...")
+        progress = gr.Progress()
         # Load models based on user selection
-        progress(0.1, desc="Loading models...")
+        progress(0, desc="Downloading models...")
         pipe = load_models(controlnet_model, diffusion_model)
-        progress(0.2, desc="Models loaded, preparing for inference...")
+        progress(0.5, desc="Models downloaded, preparing for inference...")
         
         if prompt is None or prompt == "":
             raise gr.Error("Prompt is required")
@@ -234,7 +232,7 @@ def inference(
         if count_tokens(prompt) > MAX_TOKENS:
             raise gr.Error(f"Prompt exceeds the maximum allowed tokens of {MAX_TOKENS}")
 
-        if negative_prompt and count_tokens(negative_prompt) > MAX_TOKENS:
+        if count_tokens(negative_prompt) > MAX_TOKENS:
             raise gr.Error(f"Negative prompt exceeds the maximum allowed tokens of {MAX_TOKENS}")
 
         pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)
@@ -243,7 +241,7 @@ def inference(
             seed = torch.seed()  # Generate a truly random seed
         generator = torch.manual_seed(seed)
 
-        if qr_code_content != "" or (qrcode_image is not None and qrcode_image.size == (1, 1)):
+        if qr_code_content != "" or qrcode_image.size == (1, 1):
             print("Generating QR Code from content")
             qr = qrcode.QRCode(
                 version=1,
@@ -256,11 +254,9 @@ def inference(
 
             qrcode_image = qr.make_image(fill_color=qr_color, back_color=bg_color)
             qrcode_image = resize_for_condition_image(qrcode_image, 1024)
-        elif qrcode_image is not None:
+        else:
             print("Using QR Code Image")
             qrcode_image = resize_for_condition_image(qrcode_image, 1024)
-        else:
-            raise gr.Error("No valid QR code content or image provided")
 
         # Determine which image to use as init_image and control_image
         if use_qr_code_as_init_image:
@@ -284,10 +280,10 @@ def inference(
         if invert_init_image and init_image is not None:
             init_image = invert_image(init_image)
 
-        progress(0.5, desc="Generating image...")
+        final_image = None
         out = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
+            prompt=prompt,  # Use the full prompt
+            negative_prompt=negative_prompt,  # Use the full negative prompt
             image=init_image,
             control_image=control_image,
             width=1024,
@@ -297,18 +293,19 @@ def inference(
             generator=generator,
             strength=float(strength),
             num_inference_steps=50,
-            callback=lambda step, timestep, latents: progress(0.5 + step / 100, desc=f"Step {step}/50")
         )
-        final_image = out.images[0]
+        final_image = out.images[0] if final_image is None else final_image
 
         if invert_final_image:
             final_image = invert_image(final_image)
         
-        progress(1.0, desc="Done!")
-        return final_image, seed, gr.update(visible=False)  # Add the third output here
+        return final_image, seed
+    except gr.Error as e:
+        print(f"Gradio error in inference: {str(e)}")
+        return Image.new('RGB', (1024, 1024), color='white'), -1
     except Exception as e:
         print(f"Unexpected error in inference: {str(e)}")
-        return Image.new('RGB', (1024, 1024), color='white'), -1, gr.update(visible=False)
+        return Image.new('RGB', (1024, 1024), color='white'), -1
 
 
 
@@ -479,7 +476,6 @@ with gr.Blocks(theme='Hev832/Applio', css=css, fill_width=True, fill_height=True
             
             with gr.Column():
                 result_image = gr.Image(label="Your Artistic QR Code", show_download_button=True, show_fullscreen_button=True)
-                result_image_loading = gr.Markdown("Generating image...", visible=False)
                 scan_button = gr.Button("Verify QR Code Works", visible=False)
                 scan_result = gr.Textbox(label="Validation Result of QR Code", interactive=False, visible=False)
                 used_seed = gr.Number(label="Seed Used", interactive=False)
@@ -647,119 +643,98 @@ with gr.Blocks(theme='Hev832/Applio', css=css, fill_width=True, fill_height=True
                 """
             )
 
-            def update_image_list(image, seed):
-                generated_images.value.append((image, seed))
-                return gr.Dropdown(choices=[f"Image {i+1} (Seed: {s})" for i, (_, s) in enumerate(generated_images.value)])
-
-            def load_image_for_editing(choice):
-                if not choice:
-                    return None
-                index = int(choice.split()[1]) - 1
-                return generated_images.value[index][0]
-
-            def edit_image(image, brightness, contrast, saturation):
-                if image is None:
-                    return None
-                img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
-                img = ImageEnhance.Brightness(img).enhance(brightness)
-                img = ImageEnhance.Contrast(img).enhance(contrast)
-                img = ImageEnhance.Color(img).enhance(saturation)
-                return np.array(img)
-
-            def invert_image(image):
-                if image is None:
-                    return None
-                return ImageOps.invert(Image.fromarray(image) if isinstance(image, np.ndarray) else image)
-
-            def scan_and_display(image):
-                if image is None:
-                    return "No image to scan"
-                
-                scanned_text = scan_qr_code(image)
-                if scanned_text:
-                    return f"Scanned successfully: {scanned_text}"
-                else:
-                    return "Failed to scan QR code. Try adjusting the settings for better visibility."
-
-            run_btn.click(
-                inference,
-                inputs=[
-                    qr_code_content,
-                    prompt,
-                    negative_prompt,
-                    guidance_scale,
-                    controlnet_conditioning_scale,
-                    strength,
-                    seed,
-                    init_image,
-                    qr_code_image,
-                    use_qr_code_as_init_image,
-                    sampler,
-                    bg_color,
-                    qr_color,
-                    invert_final_image,
-                    controlnet_model_dropdown,
-                    diffusion_model_dropdown,
-                    reference_image_strength,
-                ],
-                outputs=[result_image, used_seed, result_image_loading],
-                show_progress=True
-            ).then(
-                update_image_list,
-                inputs=[result_image, used_seed],
-                outputs=[image_selector]
-            )
+    def scan_and_display(image):
+        if image is None:
+            return "No image to scan"
+        
+        scanned_text = scan_qr_code(image)
+        if scanned_text:
+            return f"Scanned successfully: {scanned_text}"
+        else:
+            return "Failed to scan QR code. Try adjusting the settings for better visibility."
 
-            image_selector.change(
-                load_image_for_editing,
-                inputs=[image_selector],
-                outputs=[image_to_edit]
-            )
+    def invert_displayed_image(image):
+        if image is None:
+            return None
+        return invert_image(image)
 
-            brightness.change(
-                edit_image,
-                inputs=[image_to_edit, brightness, contrast, saturation],
-                outputs=[edited_image]
-            )
-            contrast.change(
-                edit_image,
-                inputs=[image_to_edit, brightness, contrast, saturation],
-                outputs=[edited_image]
-            )
-            saturation.change(
-                edit_image,
-                inputs=[image_to_edit, brightness, contrast, saturation],
-                outputs=[edited_image]
-            )
+    scan_button.click(
+        scan_and_display,
+        inputs=[result_image],
+        outputs=[scan_result]
+    )
 
-            invert_button.click(
-                invert_image,
-                inputs=[image_to_edit],
-                outputs=[edited_image]
-            )
+    invert_button.click(
+        invert_displayed_image,
+        inputs=[result_image],
+        outputs=[result_image]
+    )
 
-            scan_button.click(
-                scan_and_display,
-                inputs=[edited_image],
-                outputs=[scan_result]
-            )
+    invert_init_image_button.click(
+        invert_init_image_display,
+        inputs=[init_image],
+        outputs=[init_image]
+    )
+
+    brightness.change(
+        adjust_image,
+        inputs=[result_image, brightness, contrast, saturation],
+        outputs=[result_image]
+    )
+    contrast.change(
+        adjust_image,
+        inputs=[result_image, brightness, contrast, saturation],
+        outputs=[result_image]
+    )
+    saturation.change(
+        adjust_image,
+        inputs=[result_image, brightness, contrast, saturation],
+        outputs=[result_image]
+    )
+
+    # Add logic to show/hide the reference_image_strength slider
+    def update_reference_image_strength_visibility(init_image, use_qr_code_as_init_image):
+        return gr.update(visible=init_image is not None and not use_qr_code_as_init_image)
+
+    init_image.change(
+        update_reference_image_strength_visibility,
+        inputs=[init_image, use_qr_code_as_init_image],
+        outputs=[reference_image_strength]
+    )
 
-    # Define login button click behavior
-    login_button.click(
-        login,
-        inputs=[username, password],
-        outputs=[app_container, login_error, login_button]
+    use_qr_code_as_init_image.change(
+        update_reference_image_strength_visibility,
+        inputs=[init_image, use_qr_code_as_init_image],
+        outputs=[reference_image_strength]
     )
 
-    # Define password textbox submit behavior
-    password.submit(
-        login,
-        inputs=[username, password],
-        outputs=[app_container, login_error, login_button]
+    run_btn.click(
+        fn=inference,
+        inputs=[
+            qr_code_content,
+            prompt,
+            negative_prompt,
+            guidance_scale,
+            controlnet_conditioning_scale,
+            strength,
+            seed,
+            init_image,
+            qr_code_image,
+            use_qr_code_as_init_image,
+            sampler,
+            bg_color,
+            qr_color,
+            invert_final_image,
+            controlnet_model_dropdown,
+            diffusion_model_dropdown,
+            reference_image_strength,
+        ],
+        outputs=[result_image, used_seed],
+        concurrency_limit=20
     )
 
 # Load models on launch
 load_models_on_launch()
 
 blocks.queue(max_size=20)
-blocks.launch(share=False, show_api=True)
+blocks.launch(share=True, show_api=True)