File size: 1,109 Bytes
f0b80e7 24d0ffe f0b80e7 f19eb5d f0b80e7 4b79d10 f0b80e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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()
|