feat: Removes the "Pre-download Base Model" button and now automatically downloads all base models on deployment.

app.py

@@ -11,7 +11,8 @@ import requests
 import os
 import re
 import gc
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
+import time
 
 # This dummy function is required to pass the Hugging Face Spaces startup check for GPU apps.
 @spaces.GPU(duration=60)
@@ -72,6 +73,32 @@ HASH_TO_MODEL_MAP = {
     "874170688a": "RedRayz/hikari_noob_v-pred_1.2.2"
 }
 
+def download_all_base_models_on_startup():
+    """Downloads all base models listed in MODEL_LIST when the app starts."""
+    print("--- Starting pre-download of all base models ---")
+    for model_name in MODEL_LIST:
+        try:
+            print(f"Downloading: {model_name}...")
+            start_time = time.time()
+            # Handle single-file models
+            if model_name in SINGLE_FILE_MODELS:
+                filename = SINGLE_FILE_MODELS[model_name]
+                hf_hub_download(repo_id=model_name, filename=filename)
+            # Handle standard diffusers models
+            else:
+                snapshot_download(repo_id=model_name, ignore_patterns=["*.onnx", "*.flax"])
+            end_time = time.time()
+            print(f"✅ Successfully downloaded {model_name} in {end_time - start_time:.2f} seconds.")
+        except Exception as e:
+            print(f"❌ Failed to download {model_name}: {e}")
+        finally:
+            # Clean up to conserve memory
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+    print("--- Finished pre-downloading all base models ---")
+
+
 def get_civitai_file_info(version_id):
     """Gets the file metadata for a model version via the Civitai API."""
     api_url = f"https://civitai.com/api/v1/model-versions/{version_id}"
@@ -125,42 +152,6 @@ def process_long_prompt(compel_proc, prompt, negative_prompt=""):
     except Exception:
         return None, None
 
-def pre_download_base_model(model_name, progress=gr.Progress(track_tqdm=True)):
-    if not model_name:
-        return "Please select a base model to download."
-    
-    status_log = []
-    try:
-        progress(0, desc=f"Starting download for: {model_name}")
-        
-        if model_name in SINGLE_FILE_MODELS:
-            filename = SINGLE_FILE_MODELS[model_name]
-            print(f"Pre-downloading single file: {filename} from repo: {model_name}")
-            local_path = hf_hub_download(repo_id=model_name, filename=filename)
-            pipe = StableDiffusionXLPipeline.from_single_file(
-                local_path,
-                torch_dtype=torch.float16,
-                use_safetensors=True
-            )
-        else:
-            print(f"Pre-downloading diffusers model: {model_name}")
-            pipe = StableDiffusionXLPipeline.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                use_safetensors=True
-            )
-        
-        status_log.append(f"✅ Successfully downloaded {model_name}")
-        del pipe 
-    except Exception as e:
-        status_log.append(f"❌ Failed to download {model_name}: {e}")
-    finally:
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    
-    return "\n".join(status_log)
-
 def pre_download_loras(civitai_api_key, *lora_data, progress=gr.Progress(track_tqdm=True)):
     civitai_ids = lora_data[0::2]
     status_log = []
@@ -416,20 +407,21 @@ def send_info_to_txt2img(image):
     updates.append(gr.Tabs(selected=0))
     return updates
 
+# --- Execute model download on startup ---
+download_all_base_models_on_startup()
+
+
 with gr.Blocks(css="#col-container {margin: 0 auto; max-width: 1024px;}") as demo:
     gr.Markdown("# Animated SDXL T2I with LoRAs")
     with gr.Tabs(elem_id="tabs_container") as tabs:
         with gr.TabItem("txt2img", id=0):
-            gr.Markdown("<div style='background-color: #282828; color: #a0aec0; padding: 10px; border-radius: 5px; margin-bottom: 15px;'>💡 <b>Tip:</b> Pre-downloading the base model and LoRAs before clicking 'Run' can maximize your ZeroGPU time.</div>")
+            gr.Markdown("<div style='background-color: #282828; color: #a0aec0; padding: 10px; border-radius: 5px; margin-bottom: 15px;'>💡 <b>Tip:</b> Pre-downloading LoRAs before 'Run' can maximize ZeroGPU time.</div>")
             with gr.Column(elem_id="col-container"):
                 with gr.Row():
                     with gr.Column(scale=3):
                         base_model_name = gr.Dropdown(label="Base Model", choices=MODEL_LIST, value="Laxhar/noobai-XL-Vpred-1.0")
-                    with gr.Column(scale=2):
-                        with gr.Row():
-                            predownload_base_model_button = gr.Button("Pre-download Base Model")
-                            predownload_lora_button = gr.Button("Pre-download LoRAs")
-                    with gr.Column(scale=1, min_width=100):
+                    with gr.Column(scale=1):
+                        predownload_lora_button = gr.Button("Pre-download LoRAs")
                         run_button = gr.Button("Run", variant="primary")
                 
                 predownload_status = gr.Markdown("")
@@ -483,7 +475,7 @@ with gr.Blocks(css="#col-container {margin: 0 auto; max-width: 1024px;}") as dem
                 info_image_input = gr.Image(type="pil", label="Upload Image")
                 with gr.Row():
                     info_get_button = gr.Button("Get Info", variant="secondary")
-                    send_to_txt2img_button = gr.Button("Send to Txt-to-Image", variant="primary")
+                    send_to_txt2img_button = gr.Button("Send to txt2img", variant="primary")
                 gr.Markdown("### Positive Prompt"); info_prompt_output = gr.Textbox(lines=3, interactive=False, show_label=False)
                 gr.Markdown("### Negative Prompt"); info_neg_prompt_output = gr.Textbox(lines=3, interactive=False, show_label=False)
                 gr.Markdown("### Other Parameters"); info_params_output = gr.Textbox(lines=5, interactive=False, show_label=False)
@@ -500,7 +492,6 @@ with gr.Blocks(css="#col-container {margin: 0 auto; max-width: 1024px;}") as dem
 
     add_lora_button.click(fn=add_lora_row, inputs=[lora_count_state], outputs=[lora_count_state, add_lora_button] + lora_rows)
     
-    predownload_base_model_button.click(fn=pre_download_base_model, inputs=[base_model_name], outputs=[predownload_status])
     predownload_lora_button.click(fn=pre_download_loras, inputs=[civitai_api_key, *all_lora_inputs], outputs=[predownload_status])
     
     run_button.click(fn=infer,