Fix downloading CivitAI models

app.py

@@ -24,24 +24,47 @@ 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):
+    url = f"https://civitai.com/api/v1/models/{model_id}"
+    response = requests.get(url)
+    if response.status_code != 200:
+        return None
+    return response.json()
+
+def find_download_url(data, file_extension):
+    for file in data.get('modelVersions', [{}])[0].get('files', []):
+        if file['name'].endswith(file_extension):
+            return file['downloadUrl']
+    return None
+
 def download_civitai_model(model_id, lora_id=""):
     try:
-        model_url = f"https://civitai.com/api/v1/models/{model_id}"
-        model_url_download = f"https://civitai.com/api/download/models/{model_id}"
-        response = requests.get(model_url)
-        data = response.json()
-        model_name = data['name']
+        model_data = get_civitai_model_info(model_id)
+        if model_data is None:
+            return f"Error: Model with ID {model_id} not found."
+
+        model_name = model_data['name']
+        model_ckpt_url = find_download_url(model_data, '.ckpt')
+        model_safetensors_url = find_download_url(model_data, '.safetensors')
+        model_url = model_ckpt_url or model_safetensors_url
 
-        download_file(model_url_download, f"{model_name}.ckpt")
+        if not model_url:
+            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}")
 
         if lora_id:
-            lora_url = f"https://civitai.com/api/v1/models/{lora_id}"
-            lora_url_download = f"https://civitai.com/api/download/models/{lora_id}"
-            response = requests.get(lora_url)
-            lora_data = response.json()
+            lora_data = get_civitai_model_info(lora_id)
+            if lora_data is None:
+                return f"Error: LoRA with ID {lora_id} not found."
+
             lora_name = lora_data['name']
+            lora_safetensors_url = find_download_url(lora_data, '.safetensors')
+            if not lora_safetensors_url:
+                return f"Error: No suitable file found for LoRA {lora_name}."
 
-            download_file(lora_url_download, f"{lora_name}.safetensors")
+            download_file(lora_safetensors_url, f"{lora_name}.safetensors")
             loras_list.append(lora_name)
         else:
             lora_name = "None"