| import gradio as gr |
| import torch |
| from diffusers import StableDiffusionPipeline |
|
|
| model_name = "runwayml/stable-diffusion-v1-5" |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| if device == "cuda": |
| pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to(device) |
| else: |
| pipe = StableDiffusionPipeline.from_pretrained(model_name).to(device) |
|
|
| def generate_image(prompt): |
| with torch.no_grad(): |
| image = pipe(prompt, num_inference_steps=50, guidance_scale=9.5).images[0] |
|
|
| return image |
|
|
| iface = gr.Interface( |
| fn=generate_image, |
| inputs="text", |
| outputs="image", |
| title="BL's T2I Generator", |
| description="Enter a prompt to generate an image." |
| ) |
|
|
| iface.launch() |
|
|
