import gradio as gr

# Load the model
demo = gr.load("models/TDN-M/East-asian-beauty")

# Example images
examples = [
    "ea1.png",
    "ea2.png",
    "ea3.png"
]

# Create a gallery to display example images
with gr.Row():
    for example in examples:
        with gr.Column():
            gr.Gallery(value=[example], label="Example", elem_id="gallery").render() # Add elem_id

# Update the interface to include examples and CSS
demo = gr.Interface(
    fn=demo.predict,
    inputs=demo.input_components,
    outputs=demo.output_components,
    examples=examples,
    title="TDN-M/East-asian-beauty",
    description="Demo for TDN-M/East-asian-beauty model.",
    css="""
        #gallery .gallery-item {
            height: 512px;
            width: 512px;
        }
    """
)

demo.launch()