Update app.py

app.py

@@ -10,30 +10,31 @@ 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):
+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):
 
     # check if avg diff for directions need to be re-calculated
     if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]):
         clip_slider.avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1])
         x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
-        print("clip_slider.avg_diff[0]", clip_slider.avg_diff[0])
-        print("clip_slider.avg_diff[1]", clip_slider.avg_diff[1])
+
     if not sorted(slider_y) == sorted([y_concept_1, y_concept_2]):
         clip_slider.avg_diff_2nd = clip_slider.find_latent_direction(slider_y[0], slider_y[1])
         y_concept_1, y_concept_2 = slider_y[0], slider_y[1]
     
+    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=8)
     comma_concepts_x = ', '.join(slider_x)
     comma_concepts_y = ', '.join(slider_y)
 
-    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=8)
+    avg_diff_x = clip_slider.avg_diff.cpu()
+    avg_diff_y = clip_slider.avg_diff_2nd.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, 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, avg_diff_y, image
 
-def update_x(x,y,prompt):
+def update_x(x,y,prompt, avg_diff_x, avg_diff_y):
   image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
 
-def update_y(x,y,prompt):
+def update_y(x,y,prompt, avg_diff_x, avg_diff_y):
   image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
   
@@ -67,6 +68,9 @@ with gr.Blocks(css=css) as demo:
     x_concept_2 = gr.State("")
     y_concept_1 = gr.State("")
     y_concept_2 = gr.State("")
+
+    avg_diff_x = gr.State()
+    avg_diff_y = gr.State()
     
     with gr.Row():
         with gr.Column():
@@ -80,10 +84,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],
-                 outputs=[x, y, x_concept_1, x_concept_2, y_concept_1, y_concept_2, output_image])
-    x.change(fn=update_x, inputs=[x,y, prompt], outputs=[output_image])
-    y.change(fn=update_y, inputs=[x,y, prompt], outputs=[output_image])
+                 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])
 
 if __name__ == "__main__":
     demo.launch()