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

# StyleGAN2 dependencies
sys.path.append("/Users/sophiemaw/stylegan2-ada")
import dnnlib
import legacy

# Globals
stylegan_generator = None


def load_stylegan_model(
    model_path: str = "/Users/sophiemaw/Documents/VASR_NEW/pythonProject/models/generation/stylegan2-ffhq-config-f.pkl",
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
) -> None:
    global stylegan_generator

    print("⬇️ Downloading StyleGAN2 model from Hugging Face Hub...")
    model_path = hf_hub_download(
        repo_id="ZeqiangLai/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()