add steps & iterations params (#2)

- add steps & iterations params (198def9bbd7c4789dc1fe7525574bb386668499c)
- Update clip_slider_pipeline.py (1efd0cb7516df60b50fe99af918beec5b3fd29f5)

app.py

@@ -7,24 +7,24 @@ from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler,  Autoen
 #vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 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"))
 
 @spaces.GPU
-def generate(slider_x, slider_y, prompt, 
+def generate(slider_x, slider_y, prompt, iterations, steps, 
              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]):
-        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1])
+        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], iterations=iterations)
         x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
 
     if not sorted(slider_y) == sorted([y_concept_1, y_concept_2]):
-        avg_diff_2nd = clip_slider.find_latent_direction(slider_y[0], slider_y[1])
+        avg_diff_2nd = clip_slider.find_latent_direction(slider_y[0], slider_y[1], iterations=iterations)
         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, avg_diff=avg_diff, avg_diff_2nd=avg_diff_2nd)
+    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=steps, avg_diff=avg_diff, avg_diff_2nd=avg_diff_2nd)
     comma_concepts_x = ', '.join(slider_x)
     comma_concepts_y = ', '.join(slider_y)
 
@@ -36,17 +36,17 @@ def generate(slider_x, slider_y, prompt,
     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
 
 @spaces.GPU
-def update_x(x,y,prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
+def update_x(x,y,prompt, steps, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
     avg_diff = [avg_diff_x_1.cuda(), avg_diff_x_2.cuda()] 
     avg_diff_2nd = [avg_diff_y_1.cuda(), avg_diff_y_2.cuda()]
-    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
+    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
     return image
 
 @spaces.GPU
-def update_y(x,y,prompt, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
+def update_y(x,y,prompt, steps, avg_diff_x_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2):
     avg_diff = [avg_diff_x_1.cuda(), avg_diff_x_2.cuda()] 
     avg_diff_2nd = [avg_diff_y_1.cuda(), avg_diff_y_2.cuda()]
-    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=8, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
+    image = clip_slider.generate(prompt, scale=x, scale_2nd=y, num_inference_steps=steps, avg_diff=avg_diff,avg_diff_2nd=avg_diff_2nd) 
     return image
   
 css = '''
@@ -96,10 +96,12 @@ with gr.Blocks(css=css) as demo:
           y = gr.Slider(minimum=-10, value=0, maximum=10, elem_id="y", interactive=False)
           output_image = gr.Image(elem_id="image_out")
     with gr.Accordion(label="advanced options"):
+        iterations = gr.Slider(label = "num iterations", minimum=0, value=100, maximum=300)
+        steps = gr.Slider(label = "num inference steps", minimum=1, value=8, maximum=30)
         
     
     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_1, avg_diff_x_2, avg_diff_y_1, avg_diff_y_2],
+                 inputs=[slider_x, slider_y, prompt, iterations, steps, 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])

clip_slider_pipeline.py

@@ -32,17 +32,21 @@ class CLIPSlider:
 
     def find_latent_direction(self,
                               target_word:str,
-                              opposite:str):
+                              opposite:str,
+                            num_iterations: int = None):
 
         # lets identify a latent direction by taking differences between opposites
         # target_word = "happy"
         # opposite = "sad"
 
-
+        if num_iterations is not None: 
+            iterations = num_iterations
+        else:
+            iterations = self.iterations
         with torch.no_grad():
             positives = []
             negatives = []
-            for i in tqdm(range(self.iterations)):
+            for i in tqdm(range(iterations)):
                 medium = random.choice(MEDIUMS)
                 subject = random.choice(SUBJECTS)
                 pos_prompt = f"a {medium} of a {target_word} {subject}"
@@ -145,19 +149,23 @@ class CLIPSliderXL(CLIPSlider):
 
     def find_latent_direction(self,
                               target_word:str,
-                              opposite:str):
+                              opposite:str,
+                             num_iterations: int = None):
 
         # lets identify a latent direction by taking differences between opposites
         # target_word = "happy"
         # opposite = "sad"
-
+        if num_iterations is not None: 
+            iterations = num_iterations
+        else:
+            iterations = self.iterations
 
         with torch.no_grad():
             positives = []
             negatives = []
             positives2 = []
             negatives2 = []
-            for i in tqdm(range(self.iterations)):
+            for i in tqdm(range(iterations)):
                 medium = random.choice(MEDIUMS)
                 subject = random.choice(SUBJECTS)
                 pos_prompt = f"a {medium} of a {target_word} {subject}"