| |
| from diffusers import StableDiffusionPipeline |
| import torch |
| from PIL import Image |
| import gradio as gr |
|
|
| |
| model_id = "CompVis/stable-diffusion-v1-4" |
| pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) |
|
|
| |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| pipe.to(device) |
|
|
| |
| def generate_image(prompt): |
| # Generate the image |
| image = pipe(prompt).images[0] |
| |
| # Return the image |
| return image |
|
|
| |
| demo = gr.Interface( |
| fn=generate_image, |
| inputs=[gr.Textbox(label="Text Prompt", placeholder="Enter a text prompt")], |
| outputs=[gr.Image(label="Generated Image")], |
| title="Text-to-Image AI", |
| description="Enter a text prompt to generate an image", |
| ) |
|
|
| |
| demo.launch() |