import sys
import os
import torch
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# Point to your local stylegan2 utilities
#sys.path.append("src/stylegan2")  # Adjust if the path differs

import legacy

# Globals
stylegan_generator = None

def load_stylegan_model(device="cuda" if torch.cuda.is_available() else "cpu"):
    """
    Downloads and loads the StyleGAN2-FFHQ model from Hugging Face.
    """
    global stylegan_generator

    print("⬇️ Downloading StyleGAN2 model from Hugging Face Hub...")
    model_path = hf_hub_download(
        repo_id="sophiemaw/stylegan2",
        filename="stylegan2-ffhq-config-f.pkl"
    )
    print(f"🔄 Loading StyleGAN2 model from: {model_path}")
    with open(model_path, 'rb') as f:
        G = legacy.load_network_pkl(f)['G_ema'].to(device)

    stylegan_generator = G
    print("✅ StyleGAN2 model loaded and ready.")

def warm_up_models():
    """
    Loads and warms up all global models used in the pipeline.
    """
    global stylegan_generator

    if stylegan_generator is None:
        load_stylegan_model()