import requests import os import time import re # --- Configuration: Read from Environment Variables (Hugging Face Secrets) --- # Environment variables MUST be set in the Hugging Face Space Settings API_KEY = os.environ.get("CIVITAI_API_KEY") MODEL_HASH = os.environ.get("MODEL_FILE_HASH") EXTERNAL_MODEL_URL = os.environ.get("EXTERNAL_MODEL_URL") # Base Civitai API URL BASE_URL = "https://civitai.com/api/v1" # Headers for authorized requests (set dynamically inside main_upload_flow) HEADERS = {} # We extract the filename from the URL, prioritizing the 'filename=' attachment # which gives a cleaner name than relying on the URL path. match = re.search(r'filename\=\"([^\"]+)\"', EXTERNAL_MODEL_URL or "") LOCAL_FILENAME = match.group(1) if match else "uploaded_model_file.safetensors" # --- CORE FUNCTIONS (Unchanged from previous version, adapted to use global variables) --- def get_model_id_by_hash(file_hash): """ Looks up the Model Version by its SHA256 hash using the public Civitai API. If found, it returns the numeric Model ID needed to update the model. """ print(f"0. Searching for existing Model ID using hash: {file_hash[:10]}...") url = f"{BASE_URL}/model-versions/by-hash/{file_hash}" try: response = requests.get(url, headers=HEADERS) response.raise_for_status() result = response.json() model_id = result.get('modelId') if model_id: print(f" ✅ Hash match found! Model ID: {model_id}, Model Version: {result.get('id')}") return model_id else: print(" ❌ Hash lookup failed to find a matching model version.") return None except requests.exceptions.HTTPError as e: if e.response.status_code == 404: print(f" ❌ Hash lookup failed: Model version not found.") else: print(f" ❌ Error during hash lookup ({e.response.status_code}): {e}") return None except requests.exceptions.RequestException as e: print(f" ❌ Network/Request Error during hash lookup: {e}") return None def create_model_version(model_id, version_data): """ Step 1: Creates a new model version entry on Civitai and gets an upload reference. NOTE: This remains a conceptual step as the exact upload initiation endpoint for users is not part of the public documentation. """ print(f"\n1. Attempting to create new version for Model ID: {model_id}...") url = f"{BASE_URL}/models/{model_id}/versions" try: response = requests.post(url, headers=HEADERS, json=version_data) response.raise_for_status() result = response.json() print(" ✅ Version creation request sent successfully.") new_version_id = result.get('id') if not new_version_id: raise Exception("API did not return a new version ID.") print(f" New Model Version ID created: {new_version_id}") return new_version_id except requests.exceptions.RequestException as e: print(f" ❌ Error creating model version: {e}") print(f" Response Body: {response.text if 'response' in locals() else 'N/A'}") return None def download_file_from_link(url, filename): """ Step 2: Downloads the model file from the external link to the local system. """ print(f"\n2. Downloading file from external link (this may take a long time): {filename}") try: with requests.get(url, stream=True) as r: r.raise_for_status() total_size = int(r.headers.get('content-length', 0)) block_size = 1024 * 1024 # 1 Megabyte downloaded = 0 with open(filename, 'wb') as f: for data in r.iter_content(block_size): downloaded += len(data) f.write(data) # Simple progress indicator if total_size > 0: progress = downloaded / total_size * 100 print(f" {downloaded / (1024*1024):.2f} MB / {total_size / (1024*1024):.2f} MB ({progress:.1f}%)", end='\r') print(f"\n ✅ Download complete. File saved as: {filename}") return True except requests.exceptions.RequestException as e: print(f" ❌ Error downloading file. The signed URL may have expired: {e}") return False def get_presigned_upload_url(model_version_id, filename): """ Step 3: Requests a secure, time-limited upload URL from Civitai. """ print("\n3. Requesting presigned upload URL from Civitai...") url = f"{BASE_URL}/model-versions/{model_version_id}/files/upload-url" payload = { "filename": filename, "filesize": os.path.getsize(filename), "filetype": "Checkpoint", "mimetype": "application/octet-stream" } try: response = requests.post(url, headers=HEADERS, json=payload) response.raise_for_status() signed_url = response.json().get('uploadUrl') if not signed_url: raise Exception("API did not return a signed upload URL.") print(" ✅ Received secure upload URL.") return signed_url except requests.exceptions.RequestException as e: print(f" ❌ Error requesting signed URL: {e}") print(f" Response Body: {response.text if 'response' in locals() else 'N/A'}") return None def stream_upload_file(local_path, signed_url): """ Step 4: Uploads the local file directly to the presigned cloud storage URL. """ print("\n4. Streaming file content to the secure upload URL...") try: file_size = os.path.getsize(local_path) with open(local_path, 'rb') as f: upload_headers = { 'Content-Length': str(file_size), 'Content-Type': 'application/octet-stream' } response = requests.put(signed_url, data=f, headers=upload_headers) response.raise_for_status() print(" ✅ File upload successful.") return True except requests.exceptions.RequestException as e: print(f" ❌ Error during file streaming upload: {e}") print(f" Response Status: {response.status_code if 'response' in locals() else 'N/A'}") return False def main_upload_flow(): """ Executes the full simulated API upload flow, dynamically getting the model ID. """ global HEADERS # Check for required environment variables if not all([API_KEY, MODEL_HASH, EXTERNAL_MODEL_URL]): print("--- CONFIGURATION ERROR ---") if not API_KEY: print("CIVITAI_API_KEY secret is missing.") if not MODEL_HASH: print("MODEL_FILE_HASH secret is missing.") if not EXTERNAL_MODEL_URL: print("EXTERNAL_MODEL_URL secret is missing.") print("Please set these as Secrets in your Hugging Face Space settings.") return # Set up headers now that API_KEY is confirmed HEADERS.update({ "Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json" }) print("--- Civitai API Link Upload Workflow Initiated ---") # A. Get the numeric MODEL_ID using the provided hash model_id = get_model_id_by_hash(MODEL_HASH) if not model_id: print("\nFlow terminated: Could not find existing model via hash.") return # 1. Define Model Version Metadata version_metadata = { "name": "v_HF_auto_upload", # Placeholder version name "description": "Programmatic upload via Hugging Face Space using external link and hash lookup.", "baseModel": "SDXL", "trainedWords": ["dreamlookai_sdxl", "ukj_style", "hf_upload"], } # B. Create the Model Version Entry new_version_id = create_model_version(model_id, version_metadata) if not new_version_id: return # C. Download the file from the external link if not download_file_from_link(EXTERNAL_MODEL_URL, LOCAL_FILENAME): return # D. Get the secure URL for upload signed_url = get_presigned_upload_url(new_version_id, LOCAL_FILENAME) if not signed_url: return # E. Stream the downloaded file to the secure URL if not stream_upload_file(LOCAL_FILENAME, signed_url): return # F. Cleanup (Remove the downloaded file from the Hugging Face Space) try: os.remove(LOCAL_FILENAME) print(f"\n5. Cleanup complete. Model uploaded and local file '{LOCAL_FILENAME}' removed.") except OSError as e: print(f"\n5. Cleanup failed: Could not remove local file '{LOCAL_FILENAME}': {e}") print("\n**********************************************************************************") print(f"** SUCCESS: Programmatic file transfer complete for Model Version ID: {new_version_id} **") print("** Next Action: Go to the Civitai web UI to verify the file and publish the version. **") print("**********************************************************************************") if __name__ == "__main__": main_upload_flow()