Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from tqdm import tqdm
 from diffusers.models import AutoencoderKL
@@ -6,6 +7,10 @@ import gradio as gr
 import requests
 import spaces
 
+# Ensure directories exist
+os.makedirs('models', exist_ok=True)
+os.makedirs('loras', exist_ok=True)
+
 models_list = []
 loras_list = ["None"]
 models = {}
@@ -52,7 +57,7 @@ def download_and_load_civitai_model(model_id, lora_id="", progress=gr.Progress(t
             return f"Error: No suitable file found for model {model_name}."
 
         file_extension = '.ckpt' if model_ckpt_url else '.safetensors'
-        model_filename = f"{model_name}{file_extension}"
+        model_filename = os.path.join('models', f"{model_name}{file_extension}")
         download_file(model_url, model_filename)
 
         if lora_id:
@@ -65,7 +70,8 @@ def download_and_load_civitai_model(model_id, lora_id="", progress=gr.Progress(t
             if not lora_safetensors_url:
                 return f"Error: No suitable file found for LoRA {lora_name}."
 
-            download_file(lora_safetensors_url, f"{lora_name}.safetensors")
+            lora_filename = os.path.join('loras', f"{lora_name}.safetensors")
+            download_file(lora_safetensors_url, lora_filename)
             if lora_name not in loras_list:
                 loras_list.append(lora_name)
         else:
@@ -99,7 +105,8 @@ def load_model(model, lora="", use_lora=False, progress=gr.Progress(track_tqdm=T
         )
 
         if use_lora and lora != "":
-            pipeline.load_lora_weights(lora)
+            lora_path = os.path.join('loras', lora + '.safetensors')
+            pipeline.load_lora_weights(lora_path)
 
         pipeline.to("cuda")
         models[model] = pipeline