#!/usr/bin/env python3
"""
Download and cache a real Stable Diffusion model for offline use.
This runs once, then the model is cached forever.
"""

import os
import sys
from pathlib import Path

# Setup cache
cache_root = Path("d:/VSC Codes/Bild/.cache/hf")
cache_root.mkdir(parents=True, exist_ok=True)

os.environ['HF_HOME'] = str(cache_root)
os.environ['TORCH_HOME'] = str(cache_root / 'torch')

print(f"Cache directory: {cache_root}")
print()

# Use direct API instead of CLI
print("Downloading Stable Diffusion 2.1...")
print("(This is one-time setup, then cached offline)")
print()

try:
    # Use snapshot_download which is more robust
    from huggingface_hub import snapshot_download
    
    model_id = "stabilityai/stable-diffusion-2-1"
    print(f"Model: {model_id}")
    print()
    print("Downloading (this may take 5-10 minutes on first run)...")
    
    model_path = snapshot_download(
        repo_id=model_id,
        cache_dir=str(cache_root),
        local_dir_use_symlinks=False
    )
    
    print()
    print(f"✓ SUCCESS! Model cached at:")
    print(f"  {model_path}")
    print()
    print("Backend can now generate images offline!")
    sys.exit(0)
    
except Exception as e:
    print(f"✗ FAILED: {type(e).__name__}: {e}")
    print()
    print("Trying tiny-sd as fallback (smaller, faster)...")
    print()
    
    try:
        from huggingface_hub import snapshot_download
        model_id = "segmind/tiny-sd"
        model_path = snapshot_download(
            repo_id=model_id,
            cache_dir=str(cache_root),
            local_dir_use_symlinks=False
        )
        print(f"✓ Fallback successful! Cached at: {model_path}")
        sys.exit(0)
    except Exception as e2:
        print(f"✗ Fallback failed: {e2}")
        sys.exit(1)