Add civitai-to-hf-uploader

app.py

@@ -17,19 +17,20 @@ def restart_space():
     api.restart_space(repo_id="civitaiarchive/civitai-to-hf-uploader", token=os.environ["HF_TOKEN"])
 
 
-def download_file(url, file_path, folder):
+def download_file(url, file_path, folder, api_key=None):
     headers = {}
     full_path = os.path.join(folder, file_path)
     os.makedirs(os.path.dirname(full_path), exist_ok=True)
     
-    curl_cmd = ['curl', '-L', '-o', full_path, url]
+    curl_cmd = ['curl', '--fail', '-L', '-o', full_path, url]
     
     try:
         result = subprocess.run(curl_cmd, check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        if b'401' in e.stderr:
+        if b'401' in e.stderr or b'403' in e.stderr:
             # Try again with authorization
-            curl_cmd.extend(['-H', f'Authorization: Bearer {os.environ["CIVITAI_API_KEY"]}'])
+            auth_key = api_key or os.environ.get("CIVITAI_API_KEY")
+            curl_cmd.extend(['-H', f'Authorization: Bearer {auth_key}'])
             try:
                 result = subprocess.run(curl_cmd, check=True, capture_output=True, text=True)
             except subprocess.CalledProcessError as e:
@@ -87,7 +88,7 @@ def get_files_by_model_id(model_id):
     except requests.exceptions.RequestException as e:
         raise gr.Error("Something went wrong in fetching CivitAI API")
 
-def process_url(url, profile, user_repo_id, oauth_token, folder):
+def process_url(url, profile, user_repo_id, oauth_token, folder, api_key=None):
     if url.startswith("https://civitai.com/models/"):
         model_id = url.split('/')[4]
         files = get_files_by_model_id(model_id)
@@ -108,7 +109,7 @@ def process_url(url, profile, user_repo_id, oauth_token, folder):
             download_url = data['downloadUrl']
             filename = dl_path.split('/')[-1]
             gr.Info(f"Downloading {filename} ({current_file}/{total_files})")
-            download_file(download_url, dl_path, folder)
+            download_file(download_url, dl_path, folder, api_key)
             # Upload the model and card
             gr.Info(f"Uploading {filename} ({current_file}/{total_files})")
             base_folder = os.path.join(folder, os.path.dirname(dl_path))
@@ -158,7 +159,7 @@ def add_mirror(repo_id):
 
 
 
-def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuthToken, url, destination_repo):
+def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuthToken, url, destination_repo, civitai_api_key=None):
     if not profile.name:
         raise gr.Error("Are you sure you are logged in?")
     
@@ -184,7 +185,7 @@ def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuth
             gr.Info(f"Repository {user_repo_id} already exists, will update it")
         update_repo_visibility(repo_id=user_repo_id, private=False, token=oauth_token.token)
 
-        files = process_url(url, profile, user_repo_id, oauth_token.token, folder)
+        files = process_url(url, profile, user_repo_id, oauth_token.token, folder, civitai_api_key)
         if not files or len(files.keys()) == 0:
             raise gr.Error("No files were copied. Something went wrong.")
         
@@ -246,6 +247,11 @@ Once uploaded, it will add this repository to CivitaiArchive.com as a mirror.
             label="HF Repo Name",
             info="Name for the HuggingFace repository (a new one will be created if it doesn't exist)",
         )
+        civitai_api_key = gr.Textbox(
+            placeholder="Your CivitAI API key (optional)",
+            label="CivitAI API Key",
+            info="Optional: Provide your own CivitAI API key to avoid rate limits. If not provided, a default key will be used.",
+        )
         
     instructions = gr.HTML("")
     submit_button_civit = gr.Button("Upload to Hugging Face", interactive=True)
@@ -255,7 +261,7 @@ Once uploaded, it will add this repository to CivitaiArchive.com as a mirror.
     
     submit_button_civit.click(
         fn=upload_civit_to_hf,
-        inputs=[submit_source_civit, destination_repo],
+        inputs=[submit_source_civit, destination_repo, civitai_api_key],
         outputs=[output]
     )