Update app.py

app.py

@@ -40,7 +40,7 @@ qrcode_generator = qrcode.QRCode(
 
 # Define available models
 CONTROLNET_MODELS = {
-    "QR Code Monster": "monster-labs/control_v1p_sd15_qrcode_monster",
+    "QR Code Monster": "monster-labs/control_v1p_sd15_qrcode_monster/v2/",
     "QR Code": "DionTimmer/controlnet_qrcode-control_v1p_sd15",
     # Add more ControlNet models here
 }
@@ -68,7 +68,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(
@@ -76,7 +76,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
@@ -182,7 +182,7 @@ def invert_displayed_image(image):
     return inverted
 
 
-MAX_TOKENS = 100
+MAX_TOKENS = 78
 
 def count_tokens(text):
     """Count the number of tokens in the text."""
@@ -193,20 +193,21 @@ 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,
+    guidance_scale: float = 15.0,  
+    controlnet_conditioning_scale: float = 1.5,  
+    strength: float = 0.6,  
     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",
+    sampler = "DPM++ Karras SDE",  
     bg_color: str = "white",
     qr_color: str = "black",
     invert_final_image: bool = False,
     invert_init_image: bool = False,
     controlnet_model: str = "QR Code Monster",
     diffusion_model: str = "GhostMix",
+    reference_image_strength: float = 0.6,
 ):
     try:
         progress = gr.Progress()
@@ -262,19 +263,18 @@ def inference(
                 strength = 1.0
 
         # Adjust strength if using an init_image
-        if init_image is not None:
-            strength = min(strength, 0.8)  # Cap strength at 0.8 when using init_image
+        if init_image is not None and not use_qr_code_as_init_image:
+            # Map the 0-5 range to 0-1 range for the strength parameter
+            mapped_strength = min(reference_image_strength / 5.0, 1.0)
+            strength = 1.0 - mapped_strength  # Invert the strength for img2img
+        elif use_qr_code_as_init_image:
+            strength = min(strength, 0.6)  # Cap strength at 0.6 when using QR code as init_image
 
         # Invert init_image if requested
         if invert_init_image and init_image is not None:
             init_image = invert_image(init_image)
 
-        # Commenting out the prompt splitting
-        # prompt_chunks = split_prompt(prompt)
-        # negative_prompt_chunks = split_prompt(negative_prompt)
-
         final_image = None
-        # for prompt_chunk, negative_prompt_chunk in zip(prompt_chunks, negative_prompt_chunks):
         out = pipe(
             prompt=prompt,  # Use the full prompt
             negative_prompt=negative_prompt,  # Use the full negative prompt
@@ -286,7 +286,7 @@ def inference(
             controlnet_conditioning_scale=float(controlnet_conditioning_scale),
             generator=generator,
             strength=float(strength),
-            num_inference_steps=80,
+            num_inference_steps=50,
         )
         final_image = out.images[0] if final_image is None else final_image
 
@@ -301,6 +301,8 @@ def inference(
         print(f"Unexpected error in inference: {str(e)}")
         return Image.new('RGB', (1024, 1024), color='white'), -1
 
+
+
 def split_prompt(prompt, max_length=77):
     """Split the prompt into chunks that do not exceed the max_length."""
     words = prompt.split()
@@ -579,6 +581,17 @@ with gr.Blocks(theme='Hev832/Applio', css=css) as blocks:
                             """
                         )
 
+                with gr.Row():
+                    reference_image_strength = gr.Slider(
+                        minimum=0.0,
+                        maximum=5.0,
+                        step=0.05,
+                        value=0.6,
+                        label="Reference Image Influence",
+                        info="Controls how much the reference image influences the final result (0 = ignore, 5 = copy exactly)",
+                        visible=False  # We'll make this visible when a reference image is uploaded
+                    )
+
     def scan_and_display(image):
         if image is None:
             return "No image to scan"
@@ -628,6 +641,22 @@ with gr.Blocks(theme='Hev832/Applio', css=css) as blocks:
         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]
+    )
+
+    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]
+    )
+
     run_btn.click(
         inference,
         inputs=[
@@ -647,6 +676,7 @@ with gr.Blocks(theme='Hev832/Applio', css=css) as blocks:
             invert_final_image,
             controlnet_model_dropdown,
             diffusion_model_dropdown,
+            reference_image_strength,
         ],
         outputs=[result_image, used_seed],
         concurrency_limit=20
@@ -656,4 +686,4 @@ with gr.Blocks(theme='Hev832/Applio', css=css) as blocks:
 load_models_on_launch()
 
 blocks.queue(max_size=20)
-blocks.launch(share=False, show_api=True)
+blocks.launch(share=True, show_api=True)