add zerogpu dev

demo/img_gen.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import torch
 import torch.nn.functional as F
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
@@ -6,17 +7,15 @@ from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 model_id = "stabilityai/stable-diffusion-2-1"
 
-device = torch.device('cpu')
-dtype = torch.float32
-if torch.cuda.is_available():
-    device = torch.device('cuda')
-    dtype = torch.float16
+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

demo/relighting_gen.py

@@ -1,16 +1,14 @@
 import imageio
 import numpy as np
+import spaces
 import torch
 from diffusers import UniPCMultistepScheduler, StableDiffusionControlNetPipeline
 from diffusers.utils import get_class_from_dynamic_module
 
 from tqdm import tqdm
 
-device = torch.device('cpu')
-dtype = torch.float32
-if torch.cuda.is_available():
-    device = torch.device('cuda')
-    dtype = torch.float16
+device = torch.device('cuda')
+dtype = torch.float16
 
 NeuralTextureControlNetModel = get_class_from_dynamic_module(
     "dilightnet/model_helpers",
@@ -28,6 +26,7 @@ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 pipe.set_progress_bar_config(disable=True)
 
 
+@spaces.GPU
 def relighting_gen(masked_ref_img, mask, cond_path, frames, prompt, steps, seed, cfg):
     mask = mask[..., :1] / 255.
     for i in tqdm(range(frames)):