Upload zerogpu_controller.py

sd/zerogpu_controller.py

@@ -34,6 +34,8 @@ def get_first_result(img, prompt, negative_prompt,
     substrate, resized_image = sketch_process(img)
     prompt=prompt_preprocess(prompt)
 
+    FIRST_PIPE.to('cuda')
+
     result=FIRST_PIPE(image=substrate,
                             control_image=resized_image,
                             strength=strength,
@@ -44,6 +46,8 @@ def get_first_result(img, prompt, negative_prompt,
                             num_inference_steps=n_steps,
                             eta=eta)
 
+    FIRST_PIPE.to('cpu')
+
     return result.images[0]
 
 
@@ -52,6 +56,9 @@ def get_second_result(img, prompt, negative_prompt,
                       g_scale=7.5, n_steps=25, 
                       adapter_scale=0.9, adapter_factor=1.0):
 
+    DETECTOR.to('cuda')
+    SECOND_PIPE.to('cuda')
+
     preprocessed_img=DETECTOR(img, 
                               detect_resolution=1024, 
                               image_resolution=1024,
@@ -66,6 +73,9 @@ def get_second_result(img, prompt, negative_prompt,
                             adapter_conditioning_factor=adapter_factor,
                             generator = torch.manual_seed(42))
 
+    DETECTOR.to('cpu')
+    SECOND_PIPE.to('cpu')
+
     return result.images[0]