Update app.py

app.py

@@ -49,6 +49,7 @@ from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
 import torch
 import re
 import time
+import threading
 from PIL import ImageFile
 from utils import (
     download_things,
@@ -175,7 +176,11 @@ class GuiSD:
         self.last_load = datetime.now()
         self.inventory = []
 
-    def update_storage_models(self, storage_floor_gb=30, required_inventory_for_purge=3):
+        # Avoid duplicate downloads
+        self.active_downloads = set()
+        self.download_lock = threading.Lock()
+
+    def update_storage_models(self, storage_floor_gb=30, required_inventory_for_purge=4):
         while get_used_storage_gb() > storage_floor_gb:
             if len(self.inventory) < required_inventory_for_purge:
                 break
@@ -200,30 +205,48 @@ class GuiSD:
 
     def load_new_model(self, model_name, vae_model, task, controlnet_model, progress=gr.Progress(track_tqdm=True)):
 
-        # download link model > model_name
-        if model_name.startswith("http"):
-            yield f"Downloading model: {model_name}"
-            model_name = download_things(DIRECTORY_MODELS, model_name, HF_TOKEN, CIVITAI_API_KEY)
-            if not model_name:
-                raise ValueError("Error retrieving model information from URL")
+        lock_key = model_name
 
-        if IS_ZERO_GPU:
-            self.update_storage_models()
+        while True:
+            with self.download_lock:
+                if lock_key not in self.active_downloads:
+                    self.active_downloads.add(lock_key)
+                    break
 
-        vae_model = vae_model if vae_model != "None" else None
-        model_type = get_model_type(model_name)
-        dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
-
-        if not os.path.exists(model_name):
-            logger.debug(f"model_name={model_name}, vae_model={vae_model}, task={task}, controlnet_model={controlnet_model}")
-            _ = download_diffuser_repo(
-                repo_name=model_name,
-                model_type=model_type,
-                revision="main",
-                token=True,
-            )
+            yield f"Waiting for existing download to finish: {model_name}..."
+            time.sleep(1)
+
+        try:
+            # download link model > model_name
+            if model_name.startswith("http"):
+                yield f"Downloading model: {model_name}"
+                model_name = download_things(DIRECTORY_MODELS, model_name, HF_TOKEN, CIVITAI_API_KEY)
+                if not model_name:
+                    raise ValueError("Error retrieving model information from URL")
+
+            if IS_ZERO_GPU:
+                self.update_storage_models()
+
+            vae_model = vae_model if vae_model != "None" else None
+            model_type = get_model_type(model_name)
+            dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
+
+            if not os.path.exists(model_name):
+                logger.debug(f"model_name={model_name}, vae_model={vae_model}, task={task}, controlnet_model={controlnet_model}")
+                _ = download_diffuser_repo(
+                    repo_name=model_name,
+                    model_type=model_type,
+                    revision="main",
+                    token=True,
+                )
+
+            self.update_inventory(model_name)
 
-        self.update_inventory(model_name)
+        except Exception as e:
+            raise e
+        finally:
+            with self.download_lock:
+                self.active_downloads.discard(lock_key)
 
         for i in range(68):
             if not self.status_loading:
@@ -797,7 +820,7 @@ with gr.Blocks(theme=args.theme, css=CSS, fill_width=True, fill_height=False) as
                 prompt_gui = gr.Textbox(lines=5, placeholder="Enter prompt", label="Prompt")
 
                 with gr.Accordion("Negative prompt", open=False, visible=True):
-                    neg_prompt_gui = gr.Textbox(lines=3, placeholder="Enter Neg prompt", label="Negative prompt", value="bad anatomy, ((many hands, bad hands, missing fingers)), anatomical nonsense, ugly, deformed, bad proportions, bad shadow, extra limbs, missing limbs, floating limbs, disconnected limbs, malformed hands, poorly drawn, mutation, mutated hands and fingers, extra legs, interlocked fingers, extra arms, disfigured face, long neck, asymmetrical eyes, lowres, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry, duplicate, bad composition, text, worst quality, normal quality, low quality, very displeasing, monochrome, grayscale, black and white, desaturated, low contrast, muted tones, washed out, unfinished, incomplete, draft, logo, backlighting")
+                    neg_prompt_gui = gr.Textbox(lines=3, placeholder="Enter Neg prompt", label="Negative prompt", value="bad anatomy, ((many hands, bad hands, missing fingers)), anatomical nonsense, ugly, deformed, bad proportions, bad shadow, extra limbs, missing limbs, floating limbs, disconnected limbs, malformed hands, poorly drawn, mutation, mutated hands and fingers, extra legs, interlocked fingers, extra arms, disfigured face, long neck, asymmetrical eyes, lowres, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry, duplicate, bad composition, text, worst quality, normal quality, low quality, very displeasing, desaturated, low contrast, muted tones, washed out, unfinished, incomplete, draft, logo, backlighting")
                 
                 with gr.Row(equal_height=False):
                     set_params_gui = gr.Button(value="↙️", variant="secondary", size="sm")