Update app.py

app.py

@@ -14,6 +14,7 @@ import cv2
 from pyzxing import BarCodeReader
 from PIL import ImageOps
 from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import ImageEnhance
 
 from diffusers import (
     StableDiffusionPipeline,
@@ -45,31 +46,54 @@ CONTROLNET_MODELS = {
 }
 
 DIFFUSION_MODELS = {
-    "Stable Diffusion v1.5": "ckpt/sd15",
     "GhostMix": "sinkinai/GhostMix-V2-BakedVae",
-    "Real-Dream": "stablediffusionapi/edge-of-realism",
-    "Real-Dream": "stablediffusionapi/edge-of-realism",
     # Add more diffusion models here
 }
 
-def load_models(controlnet_model, diffusion_model):
-    # Download ControlNet model
-    controlnet_path = snapshot_download(CONTROLNET_MODELS[controlnet_model])
-    controlnet = ControlNetModel.from_pretrained(
+# Global variables to store loaded models
+loaded_controlnet = None
+loaded_pipe = None
+
+def load_models_on_launch():
+    global loaded_controlnet, loaded_pipe
+    print("Loading models on launch...")
+    controlnet_path = snapshot_download(CONTROLNET_MODELS["QR Code Monster"])
+    loaded_controlnet = ControlNetModel.from_pretrained(
         controlnet_path,
         torch_dtype=torch.float16
     ).to("cuda")
 
-
-    diffusion_path = snapshot_download(DIFFUSION_MODELS[diffusion_model])
-    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+    diffusion_path = snapshot_download(DIFFUSION_MODELS["GhostMix"])
+    loaded_pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
         diffusion_path,
-        controlnet=controlnet,
+        controlnet=loaded_controlnet,
         torch_dtype=torch.float16,
         safety_checker=None,
     ).to("cuda")
+    print("Models loaded successfully!")
+
+# Modify the load_models function to use global variables
+def load_models(controlnet_model, diffusion_model):
+    global loaded_controlnet, loaded_pipe
+    if loaded_controlnet is None or loaded_pipe is None:
+        load_models_on_launch()
+    return loaded_pipe
+
+# Add new functions for image adjustments
+def adjust_image(image, brightness, contrast, saturation):
+    if image is None:
+        return None
     
-    return pipe
+    img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
+    
+    if brightness != 1:
+        img = ImageEnhance.Brightness(img).enhance(brightness)
+    if contrast != 1:
+        img = ImageEnhance.Contrast(img).enhance(contrast)
+    if saturation != 1:
+        img = ImageEnhance.Color(img).enhance(saturation)
+    
+    return np.array(img)
 
 def resize_for_condition_image(input_image: Image.Image, resolution: int):
     input_image = input_image.convert("RGB")
@@ -230,13 +254,13 @@ def inference(
             negative_prompt=negative_prompt,
             image=init_image,
             control_image=control_image,
-            width=768,
-            height=768,
+            width=1024,
+            height=1024,
             guidance_scale=float(guidance_scale),
             controlnet_conditioning_scale=float(controlnet_conditioning_scale),
             generator=generator,
             strength=float(strength),
-            num_inference_steps=50,
+            num_inference_steps=100,
         )
         
         final_image = out.images[0]
@@ -247,7 +271,7 @@ def inference(
         return final_image, seed
     except Exception as e:
         print(f"Error in inference: {str(e)}")
-        return Image.new('RGB', (768, 768), color='white'), -1
+        return Image.new('RGB', (1024, 1024), color='white'), -1
 
 def invert_init_image_display(image):
     if image is None:
@@ -331,22 +355,23 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     minimum=0.0,
                     maximum=5.0,
                     step=0.01,
-                    value=1.3,
+                    value=2,
                     label="QR Code Visibility",
                 )
-                gr.Markdown(
-                    """
-                    **QR Code Visibility** determines how much the QR code itself stands out in the final design. Think of it like balancing between how "artistic" the image looks and how "functional" the QR code is.
-            
-                    - **Low settings (0.0-0.5)**: If you choose a lower value, the QR code will blend more into the art, and it might be hard to scan with a phone. This setting is great if you want the image to look amazing, but you might lose some of the scannability. Try this if you care more about art and less about the QR code being easily recognized.
-                    
-                    - **Medium settings (0.6-1.5)**: This is the sweet spot where the QR code remains clearly visible while still blending in with the art. You can still scan it easily with a phone, but it looks more creative. For most users, setting it around **1.1** is a great start to balance both art and function.
-            
-                    - **High settings (1.6-5.0)**: If you need to make sure that the QR code is super easy to scan, even if it means the image looks less like art and more like a regular QR code, then choose a higher value. This is ideal when functionality is the main goal, and the artistic side can take a backseat.
-                    
-                    Start with **1.3** if you're unsure, and adjust up or down depending on whether you want the QR code to be more artistic or more functional.
-                    """
-                )
+                with gr.Accordion("See Details", open=False):
+                    gr.Markdown(
+                        """
+                        **QR Code Visibility** determines how much the QR code itself stands out in the final design. Think of it like balancing between how "artistic" the image looks and how "functional" the QR code is.
+                
+                        - **Low settings (0.0-1)**: If you choose a lower value, the QR code will blend more into the art, and it might be hard to scan with a phone. This setting is great if you want the image to look amazing, but you might lose some of the scannability. Try this if you care more about art and less about the QR code being easily recognized.
+                        
+                        - **Medium settings (1-3)**: This is the sweet spot where the QR code remains clearly visible while still blending in with the art. You can still scan it easily with a phone, but it looks more creative. For most users, setting it around **1.1** is a great start to balance both art and function.
+                
+                        - **High settings (3-5.0)**: If you need to make sure that the QR code is super easy to scan, even if it means the image looks less like art and more like a regular QR code, then choose a higher value. This is ideal when functionality is the main goal, and the artistic side can take a backseat.
+                        
+                        Start with **1.3** if you're unsure, and adjust up or down depending on whether you want the QR code to be more artistic or more functional.
+                        """
+                    )
             
                 strength = gr.Slider(
                     minimum=0.0, 
@@ -355,17 +380,18 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     value=0.9, 
                     label="Artistic Freedom",
                 )
-                gr.Markdown(
-                    """
-                    **Artistic Freedom** controls how much the AI is allowed to change the QR code's look to match your description. It's like telling the AI how creative it can get with your QR code:
-            
-                    - **Low settings (0.0-0.3)**: If you set this low, the AI will make small changes and your QR code will look more like a regular, plain QR code. This is useful if you want something that is still creative but not too wild, keeping it simple and easy to scan.
-            
-                    - **Medium settings (0.4-0.7)**: Here, the AI will add more artistic touches but keep the QR code recognizable. You get the best of both worlds—your QR code will have some creative flair, but it will still be easy to scan. For most cases, setting it to **0.6** is a great way to keep the code functional and artistic.
-            
-                    - **High settings (0.8-1.0)**: If you set this high, the AI will go all-out creative. The QR code will look amazing, but it might be difficult to scan because the art can start to take over the code. This setting is perfect if you're aiming for a highly creative piece of art and don't mind if it's a bit harder to scan. Start at **0.9** to explore creative but functional designs.
-                    """
-                )
+                with gr.Accordion("See Details", open=False):
+                    gr.Markdown(
+                        """
+                        **Artistic Freedom** controls how much the AI is allowed to change the QR code's look to match your description. It's like telling the AI how creative it can get with your QR code:
+                
+                        - **Low settings (0.0-0.3)**: If you set this low, the AI will make small changes and your QR code will look more like a regular, plain QR code. This is useful if you want something that is still creative but not too wild, keeping it simple and easy to scan.
+                
+                        - **Medium settings (0.4-0.7)**: Here, the AI will add more artistic touches but keep the QR code recognizable. You get the best of both worlds—your QR code will have some creative flair, but it will still be easy to scan. For most cases, setting it to **0.6** is a great way to keep the code functional and artistic.
+                
+                        - **High settings (0.8-1.0)**: If you set this high, the AI will go all-out creative. The QR code will look amazing, but it might be difficult to scan because the art can start to take over the code. This setting is perfect if you're aiming for a highly creative piece of art and don't mind if it's a bit harder to scan. Start at **0.9** to explore creative but functional designs.
+                        """
+                    )
             
                 guidance_scale = gr.Slider(
                     minimum=0.0,
@@ -374,18 +400,19 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     value=7.5,
                     label="Follow the Prompt",
                 )
-                gr.Markdown(
-                    """
-                    **Follow the Prompt** tells the AI how closely it should follow your description when creating the QR code art. Think of it like giving the AI instructions on how strict or flexible it can be with your design ideas:
-            
-                    - **Low settings (0-5)**: If you choose a lower value, the AI has more freedom to get creative on its own and may not stick too closely to your description. This is great if you want to see how the AI interprets your ideas in unexpected ways.
-            
-                    - **Medium settings (5-15)**: This is a good balance where the AI will mostly follow your prompt but will also add some of its own creative touches. If you want to see some surprises but still want the design to look like what you described, start at around **7.5**.
-            
-                    - **High settings (15+)**: If you choose a higher value, the AI will stick very closely to what you wrote in the description. This is good if you have a very specific idea and don't want the AI to change much. Just keep in mind that this might limit the AI's creativity.
-                    
-                    Start at **7.5** for a balanced approach where the AI follows your ideas but still adds some artistic flair.
-                    """
+                with gr.Accordion("See Details", open=False):
+                    gr.Markdown(
+                        """
+                        **Follow the Prompt** tells the AI how closely it should follow your description when creating the QR code art. Think of it like giving the AI instructions on how strict or flexible it can be with your design ideas:
+                
+                        - **Low settings (0-5)**: If you choose a lower value, the AI has more freedom to get creative on its own and may not stick too closely to your description. This is great if you want to see how the AI interprets your ideas in unexpected ways.
+                
+                        - **Medium settings (5-15)**: This is a good balance where the AI will mostly follow your prompt but will also add some of its own creative touches. If you want to see some surprises but still want the design to look like what you described, start at around **7.5**.
+                
+                        - **High settings (15+)**: If you choose a higher value, the AI will stick very closely to what you wrote in the description. This is good if you have a very specific idea and don't want the AI to change much. Just keep in mind that this might limit the AI's creativity.
+                        
+                        Start at **7.5** for a balanced approach where the AI follows your ideas but still adds some artistic flair.
+                        """
                 )
             
                 sampler = gr.Dropdown(
@@ -393,19 +420,20 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     value="DPM++ Karras SDE", 
                     label="Art Style",
                 )
-                gr.Markdown(
-                    """
-                    **Art Style** changes how the AI creates the image, using different methods (or "samplers"). Each method has a different effect on how detailed or artistic the final QR code looks:
-            
-                    - **DPM++ Karras SDE**: This is a great all-around option for creating high-quality, detailed images. It's a good place to start if you want a balance of sharpness and creativity.
-            
-                    - **Euler**: This method creates very sharp, detailed images, making the QR code look crisp and clear. Choose this if you want a precise, well-defined design.
-            
-                    - **DDIM**: This method is better if you want the QR code to have a more artistic, abstract style. It's great for when you want the QR code to look like a piece of modern art.
-                    
-                    Feel free to experiment with different samplers to see what works best for the look you're going for!
-                    """
-                )
+                with gr.Accordion("See Details", open=False):
+                    gr.Markdown(
+                        """
+                        **Art Style** changes how the AI creates the image, using different methods (or "samplers"). Each method has a different effect on how detailed or artistic the final QR code looks:
+                
+                        - **DPM++ Karras SDE**: This is a great all-around option for creating high-quality, detailed images. It's a good place to start if you want a balance of sharpness and creativity.
+                
+                        - **Euler**: This method creates very sharp, detailed images, making the QR code look crisp and clear. Choose this if you want a precise, well-defined design.
+                
+                        - **DDIM**: This method is better if you want the QR code to have a more artistic, abstract style. It's great for when you want the QR code to look like a piece of modern art.
+                        
+                        Feel free to experiment with different samplers to see what works best for the look you're going for!
+                        """
+                    )
             
                 seed = gr.Slider(
                     minimum=-1,
@@ -415,17 +443,18 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                     label="Creative Seed",
                     randomize=False,                    
                 )
-                gr.Markdown(
-                    """
-                    **Creative Seed** controls whether the AI creates a completely new design each time or sticks to a specific design. Think of it like a recipe: with the same seed number, you get the same "recipe" for your QR code every time.
-            
-                    - **-1**: This setting makes the AI create something completely new every time you run it. Use this if you want to explore different design ideas with each attempt.
-            
-                    - **Any other number**: If you set a specific number, the AI will always create the same image based on that number. This is useful if you find a design you like and want to recreate it exactly.
-            
-                    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("See Details", open=False):
+                    gr.Markdown(
+                        """
+                        **Creative Seed** controls whether the AI creates a completely new design each time or sticks to a specific design. Think of it like a recipe: with the same seed number, you get the same "recipe" for your QR code every time.
+                
+                        - **-1**: This setting makes the AI create something completely new every time you run it. Use this if you want to explore different design ideas with each attempt.
+                
+                        - **Any other number**: If you set a specific number, the AI will always create the same image based on that number. This is useful if you find a design you like and want to recreate it exactly.
+                
+                        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", 
@@ -451,7 +480,7 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
                 )
                 diffusion_model_dropdown = gr.Dropdown(
                     choices=list(DIFFUSION_MODELS.keys()),
-                    value="Stable Diffusion v1.5",
+                    value="GhostMix",
                     label="Diffusion Model",
                     info="Select the main diffusion model for image generation"
                 )
@@ -461,9 +490,16 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
         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")
+            
+            with gr.Row():
+                brightness = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Brightness")
+                contrast = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Contrast")
+                saturation = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Saturation")
+            
             scan_result = gr.Textbox(label="Scan Result", interactive=False)
             
             gr.Markdown(
@@ -512,6 +548,22 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
         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]
+    )
+
     run_btn.click(
         inference,
         inputs=[
@@ -537,5 +589,8 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
         concurrency_limit=20
     )
 
+# Load models on launch
+load_models_on_launch()
+
 blocks.queue(max_size=20)
-blocks.launch(share=True, show_api=True)
+blocks.launch(share=False, show_api=True)