import gradio as gr
from transformers import AutoFeatureExtractor, TFAutoModel

feature_extractor = AutoFeatureExtractor.from_pretrained("akhaliq/Stylegan-ffhq-vintage")
model = TFAutoModel.from_pretrained("akhaliq/Stylegan-ffhq-vintage")

def generate_image(seed):
    inputs = feature_extractor(seed, return_tensors="pt")
    with torch.no_grad():
        images = model.generate(**inputs)
    return images[0]
    
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.inputs.Textbox(label="Input Seed for Generation"),
    outputs="image",
    title="StyleGAN3 Image Generator",
    description="Enter a seed to generate a new image with StyleGAN3."
)

iface.launch()