offloading to cpu

app.py

@@ -10,7 +10,10 @@ flash_pipe.scheduler = EulerDiscreteScheduler.from_config(flash_pipe.scheduler.c
 clip_slider = CLIPSliderXL(flash_pipe, device=torch.device("cuda"), iterations=50)
 
 @spaces.GPU
-def generate(slider_x, slider_y, prompt, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x, avg_diff_y):
+def generate(slider_x, slider_y, prompt, 
+             x_concept_1, x_concept_2, y_concept_1, y_concept_2, 
+             avg_diff_x_1, avg_diff_x_2,
+             avg_diff_y_1, avg_diff_y_2):
 
     # check if avg diff for directions need to be re-calculated
     if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]):
@@ -25,16 +28,18 @@ def generate(slider_x, slider_y, prompt, x_concept_1, x_concept_2, y_concept_1,
     comma_concepts_x = ', '.join(slider_x)
     comma_concepts_y = ', '.join(slider_y)
 
-    avg_diff_x = clip_slider.avg_diff.cpu()
-    avg_diff_y = clip_slider.avg_diff_2nd.cpu()
+    avg_diff_x_1 = clip_slider.avg_diff[0].cpu()
+    avg_diff_x_2 = clip_slider.avg_diff[1].cpu()
+    avg_diff_y_1 = clip_slider.avg_diff_2nd[0].cpu()
+    avg_diff_y_2 = clip_slider.avg_diff_2nd[1].cpu()
   
-    return gr.update(label=comma_concepts_x, interactive=True),gr.update(label=comma_concepts_y, interactive=True), x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x, avg_diff_y, image
+    return gr.update(label=comma_concepts_x, interactive=True),gr.update(label=comma_concepts_y, interactive=True), x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, image
 
-def update_x(x,y,prompt, avg_diff_x, avg_diff_y):
+def update_x(x,y,prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
   image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
 
-def update_y(x,y,prompt, avg_diff_x, avg_diff_y):
+def update_y(x,y,prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
   image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
   
@@ -69,8 +74,10 @@ with gr.Blocks(css=css) as demo:
     y_concept_1 = gr.State("")
     y_concept_2 = gr.State("")
 
-    avg_diff_x = gr.State()
-    avg_diff_y = gr.State()
+    avg_diff_x_1 = gr.State()
+    avg_diff_x_2 = gr.State()
+    avg_diff_y_1 = gr.State()
+    avg_diff_y_2 = gr.State()
     
     with gr.Row():
         with gr.Column():
@@ -84,10 +91,10 @@ with gr.Blocks(css=css) as demo:
           output_image = gr.Image(elem_id="image_out")
     
     submit.click(fn=generate,
-                 inputs=[slider_x, slider_y, prompt, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x, avg_diff_y],
-                 outputs=[x, y, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x, avg_diff_y, output_image])
-    x.change(fn=update_x, inputs=[x,y, prompt, avg_diff_x, avg_diff_y], outputs=[output_image])
-    y.change(fn=update_y, inputs=[x,y, prompt, avg_diff_x, avg_diff_y], outputs=[output_image])
+                 inputs=[slider_x, slider_y, prompt, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2],
+                 outputs=[x, y, x_concept_1, x_concept_2, y_concept_1, y_concept_2, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2, output_image])
+    x.change(fn=update_x, inputs=[x,y, prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2], outputs=[output_image])
+    y.change(fn=update_y, inputs=[x,y, prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2], outputs=[output_image])
 
 if __name__ == "__main__":
     demo.launch()