import sys
import torch
import numpy as np
from PIL import Image
import load_models

def generate_face_image(
    seed: int = None,
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
) -> Image.Image:
    """
    Generates a synthetic face using the preloaded StyleGAN2 generator.
    Args:
        seed (int): Optional seed for reproducibility.
        device (str): Device to run generation on ("cuda" or "cpu").
    Returns:
        PIL.Image.Image: Generated face image.
    """
    G = load_models.stylegan_generator

    if G is None:
        raise RuntimeError("StyleGAN generator is not loaded. Run warm_up_models() first.")

    if seed is not None:
        np.random.seed(seed)

    G = G.to(device)
    z = torch.from_numpy(np.random.randn(1, G.z_dim)).to(device)

    print("🧠 Generating face...")
    img_tensor = G(z, None, truncation_psi=0.7, noise_mode='const')
    img_array = (img_tensor.clamp(-1, 1).add(1).div(2).mul(255)
                 .permute(0, 2, 3, 1).to(torch.uint8).cpu().numpy())[0]

    return Image.fromarray(img_array, 'RGB')

# Example usage:
if __name__ == "__main__":
    img = generate_face_image()
    img.show()  # or img.save("output.jpg")