app.py

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, 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, num_images=num_images)
    elif model_choice == "DALL·E 3 XL":
        response = dalle_3.text_to_image(prompt, num_images=num_images)
    elif model_choice == "FLUX.1-dev":
        response = flux.text_to_image(prompt, num_images=num_images)
    
    # 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):
    image = generate_image(model_choice, prompt)
    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")
    ],
    outputs="image",
    title="DreamXL Image",
    description="Welcome to DreamXL Image! Choose a model and input your prompt to generate stunning visuals.",
    additional_inputs=[
        gr.Textbox(value="You are a helpful image generation assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

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