Create run_upload.py

run_upload.py

@@ -0,0 +1,244 @@
+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()
+
+