SuCicada/waifu-diffusion

1# pip install transformers gradio scipy ftfy "ipywidgets>=7,<8" datasets diffusers2import random3 4import gradio as gr5import torch6from torch import autocast7from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline8 9model_id = "hakurei/waifu-diffusion"10# torch.backends.cudnn.benchmark = True11 12device = "cuda" if torch.cuda.is_available() else "cpu"13 14 15def __def_helper():16    StableDiffusionPipeline.__call__()17 18 19pipe = StableDiffusionPipeline.from_pretrained(model_id,20                                               resume_download=True,  # 模型文件断点续传21                                               torch_dtype=torch.float16,22                                               revision='fp16')23pipe = pipe.to(device)24 25 26def infer(prompt, width, height, nums, steps, guidance_scale, seed):27    print(prompt)28    print(width, height, nums, steps, guidance_scale, seed)29 30    if prompt is not None and prompt != "":31 32        if seed is None or seed == '' or seed == -1:33            seed = int(random.randrange(4294967294))34 35        with autocast("cuda"):36            generator = torch.Generator("cuda").manual_seed(seed)37 38            images = pipe([prompt] * nums,39                          height=height,40                          width=width,41                          num_inference_steps=steps,42                          generator=generator,43                          guidance_scale=guidance_scale44                          )["sample"]45            return images46 47 48description = """49prompt 素材：[https://lexica.art](https://lexica.art) \n50seed：为空会使用随机seed51 52"""53 54 55# with block as demo:56def run():57    _app = gr.Interface(58        fn=infer,59        title="Waifu Diffusion",60        description=description,61        inputs=[62            gr.Textbox(label="输入 prompt"),63            gr.Slider(512, 1024, 512, step=64, label="width"),64            gr.Slider(512, 1024, 512, step=64, label="height"),65            gr.Slider(1, 4, 1, step=1, label="Number of Images"),66            gr.Slider(10, 150, step=1, value=50,67                      label="num_inference_steps:\n"68                            "去噪步骤的数量。更多的去噪步骤通常会导致更高质量的图像，但会降低推理速度。"),69            gr.Slider(0, 20, 7.5, step=0.5,70                      label="guidance_scale:\n" +71                            "较高的引导比例鼓励生成与文本“提示”密切相关的图像，通常以降低图像质量为代价"),72            gr.Textbox(label="随机 seed",73                       placeholder="Random Seed",74                       lines=1),75        ],76        outputs=[77            gr.Gallery(label="Generated images")78        ])79 80    return _app81 82 83app = run()84app.launch(debug=True)85