imageToText / app.py
Alpha108's picture
Create app.py
eeaed51 verified
raw
history blame contribute delete
782 Bytes
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
# Load the Stable Diffusion model from Hugging Face Hub
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
def generate_image(prompt):
image = pipe(prompt).images[0]
return image
# Gradio UI
demo = gr.Interface(
fn=generate_image,
inputs=gr.Textbox(label="Enter your prompt"),
outputs=gr.Image(label="Generated Image"),
title="Text-to-Image Generator",
description="Generate images from text using Stable Diffusion (v1-5)."
)
if __name__ == "__main__":
demo.launch()