app.py

import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

# List of models for text-to-image generation
model_names = [
    "runwayml/stable-diffusion-v1-5",
    "stabilityai/stable-diffusion-2-1",
    "CompVis/stable-diffusion-v1-4"
]

# Placeholder for pipeline
pipe = None

# Function to load model
def load_model(model_name):
    global pipe
    try:
        pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
        pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
        return f"Model {model_name} loaded successfully!"
    except Exception as e:
        return f"Failed to load model: {str(e)}"

# Function to generate image from text
def generate_image(prompt, guidance_scale):
    if pipe is None:
        return None, "Model not loaded yet!"
    try:
        image = pipe(prompt, guidance_scale=guidance_scale).images[0]
        return image, "Image generated successfully!"
    except Exception as e:
        return None, f"Image generation failed: {str(e)}"

# Gradio UI
with gr.Blocks() as demo:
    with gr.Row():
        model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="runwayml/stable-diffusion-v1-5")
        load_button = gr.Button("Load Model")
    load_status = gr.Textbox(label="Status", interactive=False)

    with gr.Row():
        prompt_input = gr.Textbox(label="Enter text prompt", placeholder="A futuristic city in the sky", lines=2)
        guidance_slider = gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale")
    
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")
    message = gr.Textbox(label="Generation Status", interactive=False)

    load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_status)
    generate_button.click(fn=generate_image, inputs=[prompt_input, guidance_slider], outputs=[output_image, message])

demo.launch()