import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

# Create InferenceClients for each model (you may need to adjust based on model endpoints)
stable_diffusion = InferenceClient("stabilityai/stable-diffusion-3.5-large-turbo")
dalle_3 = InferenceClient("ehristoforu/dalle-3-xl-v2")
flux = InferenceClient("black-forest-labs/FLUX.1-dev")

def generate_image(model_choice, prompt, negative_prompt, image_size, num_images=1):
    """Function to generate images based on the chosen model."""
    if model_choice == "Stable Diffusion 3.5 Large Turbo":
        response = stable_diffusion.text_to_image(prompt, negative_prompt=negative_prompt, num_images=num_images, size=image_size)
    elif model_choice == "DALL·E 3 XL":
        response = dalle_3.text_to_image(prompt, negative_prompt=negative_prompt, num_images=num_images, size=image_size)
    elif model_choice == "FLUX.1-dev":
        response = flux.text_to_image(prompt, negative_prompt=negative_prompt, num_images=num_images, size=image_size)
    
    # Return the generated images (assuming each model returns a URL or image object)
    return response[0]["image"]  # Adjust as needed based on actual response format

# Create a function to handle user input
def generate_image_response(prompt, model_choice, negative_prompt, image_size, num_images):
    image = generate_image(model_choice, prompt, negative_prompt, image_size, num_images)
    return image

# Define Gradio Interface
demo = gr.Interface(
    fn=generate_image_response, 
    inputs=[
        gr.Textbox(label="Enter your prompt here"),
        gr.Dropdown(choices=["Stable Diffusion 3.5 Large Turbo", "DALL·E 3 XL", "FLUX.1-dev"], label="Choose Model", value="Stable Diffusion 3.5 Large Turbo"),
        gr.Textbox(value="", label="Negative prompt (what you don't want in the image)"),
        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of images to generate"),
        gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Image size (width & height)"),
    ],
    outputs="image",
    title="DreamXL Image",
    description="Welcome to DreamXL Image! Choose a model, input your prompt, negative prompt, and set image parameters to generate stunning visuals.",
)

if __name__ == "__main__":
    demo.launch()