Update app.py

app.py

@@ -13,6 +13,7 @@ import numpy as np
 import cv2
 from pyzxing import BarCodeReader
 from PIL import ImageOps
+from huggingface_hub import hf_hub_download, snapshot_download
 
 from diffusers import (
     StableDiffusionPipeline,
@@ -51,14 +52,17 @@ DIFFUSION_MODELS = {
 }
 
 def load_models(controlnet_model, diffusion_model):
-    
+    # Download ControlNet model
+    controlnet_path = snapshot_download(CONTROLNET_MODELS[controlnet_model])
     controlnet = ControlNetModel.from_pretrained(
-        CONTROLNET_MODELS[controlnet_model],
+        controlnet_path,
         torch_dtype=torch.float16
     ).to("cuda")
 
+
+    diffusion_path = snapshot_download(DIFFUSION_MODELS[diffusion_model])
     pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-        DIFFUSION_MODELS[diffusion_model],
+        diffusion_path,
         controlnet=controlnet,
         torch_dtype=torch.float16,
         safety_checker=None,
@@ -145,7 +149,7 @@ def invert_displayed_image(image):
         return Image.fromarray(inverted)
     return inverted
 
-#@spaces.GPU()
+@spaces.GPU()
 def inference(
     qr_code_content: str,
     prompt: str,
@@ -164,10 +168,13 @@ def inference(
     invert_init_image: bool = False,  # New parameter
     controlnet_model: str = "QR Code Monster",
     diffusion_model: str = "Stable Diffusion v1.5",
+    progress=gr.Progress()
 ):
     try:
         # Load models based on user selection
+        progress(0, desc="Downloading models...")
         pipe = load_models(controlnet_model, diffusion_model)
+        progress(0.5, desc="Models downloaded, preparing for inference...")
         
         if prompt is None or prompt == "":
             raise gr.Error("Prompt is required")
@@ -524,10 +531,11 @@ with gr.Blocks(theme='Hev832/Applio') as blocks:
             invert_init_image,
             controlnet_model_dropdown,
             diffusion_model_dropdown,
+            gr.Progress()
         ],
         outputs=[result_image, used_seed],
         concurrency_limit=20
     )
 
 blocks.queue(max_size=20)
-blocks.launch(share=False, show_api=False)
+blocks.launch(share=True, show_api=True)