| import os |
| import gradio as gr |
| import torch |
| from diffusers import StableDiffusionPipeline |
|
|
| hf_token = os.getenv("HF_TOKEN") |
|
|
| |
| model_id = "runwayml/stable-diffusion-v1-5" |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| dtype = torch.float16 if device == "cuda" else torch.float32 |
|
|
| |
| pipe = StableDiffusionPipeline.from_pretrained( |
| model_id, |
| dtype=dtype, |
| token=hf_token, |
| safety_checker=None |
| ).to(device) |
|
|
| |
| def generate_image(prompt): |
| image = pipe(prompt).images[0] |
| return image |
|
|
| |
| demo = gr.Interface( |
| fn=generate_image, |
| inputs="text", |
| outputs="image", |
| title="SD1.5 文生图 Demo" |
| ) |
|
|
| demo.launch() |
|
|
