app.py

import gradio as gr
from transformers import pipeline
import spaces

# Load the pipeline (lazy-load the model to save resources)
@spaces.GPU
def load_model():
    return pipeline(
        "text-to-image",
        model="stabilityai/stable-diffusion-2-1",
        torch_dtype="float16"  # Ensure compatibility with ZeroGPU
    )

# Initialize the pipeline
model = load_model()

# Function to generate images
@spaces.GPU
def generate_image(prompt, guidance_scale=7.5):
    print(f"Generating image for prompt: {prompt}")
    images = model(prompt, guidance_scale=guidance_scale)
    return images[0]

# Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe your image..."),
        gr.Slider(1, 20, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Text-to-Image Generator",
    description="Generate images from text prompts using Stable Diffusion."
)

# Launch the app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)