Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,036 Bytes
084ab29 af6bfff 084ab29 af6bfff 084ab29 af6bfff 084ab29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import gradio as gr
import spaces
import torch
import torch.nn.functional as F
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
model_id = "stabilityai/stable-diffusion-2-1"
device = torch.device('cuda')
dtype = torch.float16
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)
@spaces.GPU
def img_gen(prompt, seed, steps, cfg, down_from_768=False, progress=gr.Progress(track_tqdm=True)):
generator = torch.Generator(device=device).manual_seed(int(seed))
hw = 512 if not down_from_768 else 768
image = pipe(prompt, generator=generator, num_inference_steps=int(steps), guidance_scale=cfg, output_type='np', height=hw, width=hw).images[0]
if down_from_768:
image = F.interpolate(torch.from_numpy(image)[None].permute(0, 3, 1, 2), size=(512, 512), mode='bilinear', align_corners=False, antialias=True).permute(0, 2, 3, 1)[0].cpu().numpy()
return image
|