Update app.py

app.py

@@ -10,29 +10,30 @@ flash_pipe.scheduler = EulerDiscreteScheduler.from_config(flash_pipe.scheduler.c
 clip_slider = CLIPSliderXL(flash_pipe, device=torch.device("cuda"), iterations=150)
 
 @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])
+        avg_diff_x = clip_slider.find_latent_direction(slider_x[0], slider_x[1])
         x_concept_1, x_concept_2 = slider_x[0], slider_x[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])
+        avg_diff_y = clip_slider.find_latent_direction(slider_y[0], slider_y[1])
         y_concept_1, y_concept_2 = slider_y[0], slider_y[1]
     
     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)
+    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=8, avg_diff=avg_diff_x, avg_diff_2nd=avg_diff_y)
   
-    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):
-  image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
+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, avg_diff=avg_diff_x, avg_diff_2nd=avg_diff_y) 
   return image
 
-def update_y(x,y,prompt):
-  image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
+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, avg_diff=avg_diff_x, avg_diff_2nd=avg_diff_y) 
   return image
   
 css = '''
@@ -65,6 +66,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():
@@ -78,10 +82,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()