run_upload.py

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()