env

app.py

@@ -7,13 +7,18 @@ from extension import CustomStableDiffusionControlNetPipeline
 
 negative_prompt = ""
 device = torch.device('cuda')
-controlnet = ControlNetModel.from_pretrained("BlockDetail/PartialSketchControlNet", torch_dtype=torch.float16).to(device)
-pipe = CustomStableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    controlnet=controlnet, torch_dtype=torch.float16
-).to(device)
-pipe.safety_checker = None
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe = None
+
++@spaces.GPU
+def load():
+    global pipe
+    controlnet = ControlNetModel.from_pretrained("BlockDetail/PartialSketchControlNet", torch_dtype=torch.float16).to(device)
+    pipe = CustomStableDiffusionControlNetPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        controlnet=controlnet, torch_dtype=torch.float16
+    ).to(device)
+    pipe.safety_checker = None
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 threshold = 250
 curr_num_samples = 2
 
@@ -48,9 +53,11 @@ with gr.Blocks() as demo:
         start_state.append([None, None])
     sketch_states = gr.State(start_state)
     checkbox_state = gr.State(True)
-        
+    
+    +@spaces.GPU
     def sketch(curr_sketch_image, dilation_mask, prompt, seed, num_steps, dilation):
         global curr_num_samples
+        global pipe
         generator = torch.Generator(device="cuda:0")
         generator.manual_seed(seed)