Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from diffusers import StableDiffusionPipeline | |
| import torch | |
| # Automatically detect device | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Load Stable Diffusion model | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-2", | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32 | |
| ).to(device) | |
| # Define image generation function | |
| def generate_image(prompt): | |
| image = pipe(prompt).images[0] | |
| return image | |
| # Gradio Interface | |
| gr.Interface( | |
| fn=generate_image, | |
| inputs=gr.Textbox(label="Enter your prompt here"), | |
| outputs="image", | |
| title="Text to Image Generator with Stable Diffusion 2" | |
| ).launch() | |