Enhanced the download and load functions

README.md

@@ -4,7 +4,6 @@ emoji: 🎴
 colorFrom: red
 colorTo: pink
 sdk: gradio
-sdk_version: 4.24.0
 app_file: app.py
 pinned: true
 license: apache-2.0

app.py

@@ -24,7 +24,7 @@ def download_file(url, filename, progress=gr.Progress(track_tqdm=True)):
     if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
         print("ERROR, something went wrong")
 
-def get_civitai_model_info(model_id, progress=gr.Progress(track_tqdm=True)):
+def get_civitai_model_info(model_id):
     url = f"https://civitai.com/api/v1/models/{model_id}"
     response = requests.get(url)
     if response.status_code != 200:
@@ -37,7 +37,7 @@ def find_download_url(data, file_extension):
             return file['downloadUrl']
     return None
 
-def download_civitai_model(model_id, lora_id="", progress=gr.Progress(track_tqdm=True)):
+def download_and_load_civitai_model(model_id, lora_id="", progress=gr.Progress(track_tqdm=True)):
     try:
         model_data = get_civitai_model_info(model_id)
         if model_data is None:
@@ -52,7 +52,8 @@ def download_civitai_model(model_id, lora_id="", progress=gr.Progress(track_tqdm
             return f"Error: No suitable file found for model {model_name}."
 
         file_extension = '.ckpt' if model_ckpt_url else '.safetensors'
-        download_file(model_url, f"{model_name}{file_extension}")
+        model_filename = f"{model_name}{file_extension}"
+        download_file(model_url, model_filename)
 
         if lora_id:
             lora_data = get_civitai_model_info(lora_id)
@@ -65,12 +66,17 @@ def download_civitai_model(model_id, lora_id="", progress=gr.Progress(track_tqdm
                 return f"Error: No suitable file found for LoRA {lora_name}."
 
             download_file(lora_safetensors_url, f"{lora_name}.safetensors")
-            loras_list.append(lora_name)
+            if lora_name not in loras_list:
+                loras_list.append(lora_name)
         else:
             lora_name = "None"
 
-        models_list.append(model_name)
-        return "Model/LoRA Downloaded!"
+        if model_name not in models_list:
+            models_list.append(model_name)
+
+        # Load model after downloading
+        load_result = load_model(model_filename, lora_name, use_lora=(lora_name != "None"))
+        return f"Model/LoRA Downloaded and Loaded! {load_result}"
     except Exception as e:
         return f"Error downloading model or LoRA: {e}"
 
@@ -114,11 +120,6 @@ def generate_images(
     progress=gr.Progress(track_tqdm=True)
 ):
     if prompt is not None and prompt.strip() != "":
-        if lora_name == "None":
-            load_model(model_name, "", False)
-        elif lora_name in loras_list and lora_name != "None":
-            load_model(model_name, lora_name, True)
-
         pipe = models.get(model_name)
         if pipe is None:
             return []
@@ -172,6 +173,6 @@ with gr.Blocks(theme='ParityError/Interstellar') as demo:
 
             download_output = gr.Textbox(label="Download Output")
 
-            download_button.click(download_civitai_model, inputs=[model_id, lora_id], outputs=download_output)
+            download_button.click(download_and_load_civitai_model, inputs=[model_id, lora_id], outputs=download_output)
 
-demo.launch()
+    demo.launch()