revert taking out avg_diff from attributes part 2

app.py

@@ -7,34 +7,34 @@ from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
 
 flash_pipe = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash").to("cuda", torch.float16)
 flash_pipe.scheduler = EulerDiscreteScheduler.from_config(flash_pipe.scheduler.config)
-clip_slider = CLIPSliderXL(flash_pipe, device=torch.device("cuda"), iterations=100)
+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,  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):
 
     # check if avg diff for directions need to be re-calculated
     if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]):
-        avg_diff_x = clip_slider.find_latent_direction(slider_x[0], slider_x[1])
+        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]):
-        avg_diff_y = clip_slider.find_latent_direction(slider_y[0], slider_y[1])
+        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]
     
     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=avg_diff_x, avg_diff_2nd=avg_diff_y)
+    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=8)
   
-    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, image
 
-@spaces.GPU
-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) 
+def update_x(x,y,prompt):
+  image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
 
-@spaces.GPU
-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) 
+def update_y(x,y,prompt):
+  image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8) 
   return image
   
 css = '''
@@ -67,9 +67,6 @@ 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(None)
-    avg_diff_y = gr.State(None)
     
     with gr.Row():
         with gr.Column():
@@ -83,10 +80,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],
+                 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])
 
 if __name__ == "__main__":
     demo.launch()