import gradio as gr
from pipeline import SmileGen
import torch
from PIL import Image
import numpy as np
import os


def read_samples(path):
    # read the samples from the path
    samples = []
    for filename in os.listdir(path):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            img = Image.open(os.path.join(path, filename))
            samples.append(np.array(img))
    return samples

def create_image_generation_demo():
    # load sample images
    image_list = []

    model = SmileGen()

    demo = gr.Interface(
        fn=model.run,
        inputs=[
            gr.Image(label="Input Image", type="pil")
        ],
        outputs=[
            gr.Image(label="Generated Image")
        ],
        title="Smile!",
        description="Upload an image and generate a new image using a custom pipeline.",
        examples=image_list
    )
    
    return demo

# Launch the demo
if __name__ == "__main__":
    demo = create_image_generation_demo()
    demo.launch()