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