env

app.py

@@ -4,6 +4,7 @@ import torch
 import numpy as np
 from PIL import Image, ImageFilter
 from extension import CustomStableDiffusionControlNetPipeline
+import spaces
 
 negative_prompt = ""
 device = torch.device('cuda')
@@ -11,6 +12,7 @@ pipe = None
 
 print(gr.__version__)
 
+@spaces.GPU
 def load():
     global pipe
     controlnet = ControlNetModel.from_pretrained("BlockDetail/PartialSketchControlNet", torch_dtype=torch.float16).to(device)
@@ -54,6 +56,7 @@ with gr.Blocks() as demo:
     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
@@ -159,5 +162,5 @@ with gr.Blocks() as demo:
     stroke_type[0].change(change_color, [stroke_type[0]], canvas)
     num_samples[0].change(change_num_samples, [num_samples[0]], None)
 
-
+load()
 demo.launch(share = True, debug = True)