take avg_diff out of attributes and save to cpu (#1)

- take avg_diff out of attributes and save to cpu (a64b549e3ec96a9cf832c237af6bd7dd3baae807)
- Update clip_slider_pipeline.py (224212aeaf4f1ed103d432f5d08e8eee69601810)

app.py

@@ -17,14 +17,14 @@ def generate(slider_x, slider_y, prompt,
 
     # 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 = 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_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)
+    image = clip_slider.generate(prompt, scale=0, scale_2nd=0, num_inference_steps=8, avg_diff=avg_diff, avg_diff_2nd=avg_diff_2nd)
     comma_concepts_x = ', '.join(slider_x)
     comma_concepts_y = ', '.join(slider_y)
 
@@ -36,11 +36,15 @@ 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
 
 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) 
+    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) 
   return image
 
 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) 
+    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) 
   return image
   
 css = '''

clip_slider_pipeline.py

@@ -73,6 +73,8 @@ class CLIPSlider:
         only_pooler = False,
         normalize_scales = False, # whether to normalize the scales when avg_diff_2nd is not None
         correlation_weight_factor = 1.0,
+        avg_diff = None,
+        avg_diff_2nd = None, 
         **pipeline_kwargs
         ):
         # if doing full sequence, [-0.3,0.3] work well, higher if correlation weighted is true
@@ -83,14 +85,14 @@ class CLIPSlider:
                                   max_length=self.pipe.tokenizer.model_max_length).input_ids.cuda()
         prompt_embeds = self.pipe.text_encoder(toks).last_hidden_state
 
-        if self.avg_diff_2nd and normalize_scales:
+        if avg_diff_2nd and normalize_scales:
             denominator = abs(scale) + abs(scale_2nd)
             scale = scale / denominator
             scale_2nd = scale_2nd / denominator
         if only_pooler:
-            prompt_embeds[:, toks.argmax()] = prompt_embeds[:, toks.argmax()] + self.avg_diff * scale
-            if self.avg_diff_2nd:
-                prompt_embeds[:, toks.argmax()] += self.avg_diff_2nd * scale_2nd
+            prompt_embeds[:, toks.argmax()] = prompt_embeds[:, toks.argmax()] + avg_diff * scale
+            if avg_diff_2nd:
+                prompt_embeds[:, toks.argmax()] += avg_diff_2nd * scale_2nd
         else:
             normed_prompt_embeds = prompt_embeds / prompt_embeds.norm(dim=-1, keepdim=True)
         sims = normed_prompt_embeds[0] @ normed_prompt_embeds[0].T
@@ -102,9 +104,9 @@ class CLIPSlider:
 
         # weights = torch.sigmoid((weights-0.5)*7)
         prompt_embeds = prompt_embeds + (
-                    weights * self.avg_diff[None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale)
-        if self.avg_diff_2nd:
-            prompt_embeds += weights * self.avg_diff_2nd[None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale_2nd
+                    weights * avg_diff[None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale)
+        if avg_diff_2nd:
+            prompt_embeds += weights * avg_diff_2nd[None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale_2nd
 
 
         torch.manual_seed(seed)
@@ -198,6 +200,8 @@ class CLIPSliderXL(CLIPSlider):
         only_pooler = False,
         normalize_scales = False,
         correlation_weight_factor = 1.0,
+        avg_diff = None,
+        avg_diff_2nd = None, 
         **pipeline_kwargs
         ):
         # if doing full sequence, [-0.3,0.3] work well, higher if correlation weighted is true
@@ -232,16 +236,16 @@ class CLIPSliderXL(CLIPSlider):
                 pooled_prompt_embeds = prompt_embeds[0]
                 prompt_embeds = prompt_embeds.hidden_states[-2]
 
-                if self.avg_diff_2nd and normalize_scales:
+                if avg_diff_2nd and normalize_scales:
                     denominator = abs(scale) + abs(scale_2nd)
                     scale = scale / denominator
                     scale_2nd = scale_2nd / denominator
                 if only_pooler:
-                    prompt_embeds[:, toks.argmax()] = prompt_embeds[:, toks.argmax()] + self.avg_diff[0] * scale
-                    if self.avg_diff_2nd:
-                        prompt_embeds[:, toks.argmax()] += self.avg_diff_2nd[0] * scale_2nd
+                    prompt_embeds[:, toks.argmax()] = prompt_embeds[:, toks.argmax()] + avg_diff[0] * scale
+                    if avg_diff_2nd:
+                        prompt_embeds[:, toks.argmax()] += avg_diff_2nd[0] * scale_2nd
                 else:
-                    print(self.avg_diff)
+                    print(avg_diff)
                     normed_prompt_embeds = prompt_embeds / prompt_embeds.norm(dim=-1, keepdim=True)
                     sims = normed_prompt_embeds[0] @ normed_prompt_embeds[0].T
 
@@ -251,18 +255,18 @@ class CLIPSliderXL(CLIPSlider):
                         standard_weights = torch.ones_like(weights)
 
                         weights = standard_weights + (weights - standard_weights) * correlation_weight_factor
-                        prompt_embeds = prompt_embeds + (weights * self.avg_diff[0][None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale)
-                        if self.avg_diff_2nd:
-                            prompt_embeds += (weights * self.avg_diff_2nd[0][None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale_2nd)
+                        prompt_embeds = prompt_embeds + (weights * avg_diff[0][None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale)
+                        if avg_diff_2nd:
+                            prompt_embeds += (weights * avg_diff_2nd[0][None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale_2nd)
                     else:
                         weights = sims[toks.argmax(), :][None, :, None].repeat(1, 1, 1280)
 
                         standard_weights = torch.ones_like(weights)
 
                         weights = standard_weights + (weights - standard_weights) * correlation_weight_factor
-                        prompt_embeds = prompt_embeds + (weights * self.avg_diff[1][None, :].repeat(1, self.pipe.tokenizer_2.model_max_length, 1) * scale)
-                        if self.avg_diff_2nd:
-                            prompt_embeds += (weights * self.avg_diff_2nd[1][None, :].repeat(1, self.pipe.tokenizer_2.model_max_length, 1) * scale_2nd)
+                        prompt_embeds = prompt_embeds + (weights * avg_diff[1][None, :].repeat(1, self.pipe.tokenizer_2.model_max_length, 1) * scale)
+                        if avg_diff_2nd:
+                            prompt_embeds += (weights * avg_diff_2nd[1][None, :].repeat(1, self.pipe.tokenizer_2.model_max_length, 1) * scale_2nd)
 
                 bs_embed, seq_len, _ = prompt_embeds.shape
                 prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
@@ -276,150 +280,3 @@ class CLIPSliderXL(CLIPSlider):
                          **pipeline_kwargs).images[0]
 
         return image
-
-
-class CLIPSlider3(CLIPSlider):
-    def find_latent_direction(self,
-                              target_word:str,
-                              opposite:str):
-
-        # lets identify a latent direction by taking differences between opposites
-        # target_word = "happy"
-        # opposite = "sad"
-
-
-        with torch.no_grad():
-            positives = []
-            negatives = []
-            positives2 = []
-            negatives2 = []
-            for i in tqdm(range(self.iterations)):
-                medium = random.choice(MEDIUMS)
-                subject = random.choice(SUBJECTS)
-                pos_prompt = f"a {medium} of a {target_word} {subject}"
-                neg_prompt = f"a {medium} of a {opposite} {subject}"
-
-                pos_toks = self.pipe.tokenizer(pos_prompt, return_tensors="pt", padding="max_length", truncation=True,
-                                          max_length=self.pipe.tokenizer.model_max_length).input_ids.cuda()
-                neg_toks = self.pipe.tokenizer(neg_prompt, return_tensors="pt", padding="max_length", truncation=True,
-                                          max_length=self.pipe.tokenizer.model_max_length).input_ids.cuda()
-                pos = self.pipe.text_encoder(pos_toks).text_embeds
-                neg = self.pipe.text_encoder(neg_toks).text_embeds
-                positives.append(pos)
-                negatives.append(neg)
-
-                pos_toks2 = self.pipe.tokenizer_2(pos_prompt, return_tensors="pt", padding="max_length", truncation=True,
-                                             max_length=self.pipe.tokenizer_2.model_max_length).input_ids.cuda()
-                neg_toks2 = self.pipe.tokenizer_2(neg_prompt, return_tensors="pt", padding="max_length", truncation=True,
-                                             max_length=self.pipe.tokenizer_2.model_max_length).input_ids.cuda()
-                pos2 = self.pipe.text_encoder_2(pos_toks2).text_embeds
-                neg2 = self.pipe.text_encoder_2(neg_toks2).text_embeds
-                positives2.append(pos2)
-                negatives2.append(neg2)
-
-        positives = torch.cat(positives, dim=0)
-        negatives = torch.cat(negatives, dim=0)
-        diffs = positives - negatives
-        avg_diff = diffs.mean(0, keepdim=True)
-
-        positives2 = torch.cat(positives2, dim=0)
-        negatives2 = torch.cat(negatives2, dim=0)
-        diffs2 = positives2 - negatives2
-        avg_diff2 = diffs2.mean(0, keepdim=True)
-        return (avg_diff, avg_diff2)
-
-    def generate(self,
-        prompt = "a photo of a house",
-        scale = 2,
-        seed = 15,
-        only_pooler = False,
-        correlation_weight_factor = 1.0,
-        ** pipeline_kwargs
-        ):
-        # if doing full sequence, [-0.3,0.3] work well, higher if correlation weighted is true
-        # if pooler token only [-4,4] work well
-        clip_text_encoders = [self.pipe.text_encoder, self.pipe.text_encoder_2]
-        clip_tokenizers = [self.pipe.tokenizer, self.pipe.tokenizer_2]
-        with torch.no_grad():
-            # toks = pipe.tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=77).input_ids.cuda()
-            # prompt_embeds = pipe.text_encoder(toks).last_hidden_state
-
-            clip_prompt_embeds_list = []
-            clip_pooled_prompt_embeds_list = []
-            for i, text_encoder in enumerate(clip_text_encoders):
-
-                if i < 2:
-                    tokenizer = clip_tokenizers[i]
-                    text_inputs = tokenizer(
-                        prompt,
-                        padding="max_length",
-                        max_length=tokenizer.model_max_length,
-                        truncation=True,
-                        return_tensors="pt",
-                    )
-                    toks = text_inputs.input_ids
-
-                    prompt_embeds = text_encoder(
-                        toks.to(text_encoder.device),
-                        output_hidden_states=True,
-                    )
-
-                    # We are only ALWAYS interested in the pooled output of the final text encoder
-                    pooled_prompt_embeds = prompt_embeds[0]
-                    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
-                    clip_pooled_prompt_embeds_list.append(pooled_prompt_embeds)
-                    prompt_embeds = prompt_embeds.hidden_states[-2]
-                else:
-                    text_inputs = self.pipe.tokenizer_3(
-                        prompt,
-                        padding="max_length",
-                        max_length=self.tokenizer_max_length,
-                        truncation=True,
-                        add_special_tokens=True,
-                        return_tensors="pt",
-                    )
-                    toks = text_inputs.input_ids
-                    prompt_embeds = self.pipe.text_encoder_3(toks.to(self.device))[0]
-                    t5_prompt_embed_shape = prompt_embeds.shape[-1]
-
-                if only_pooler:
-                    prompt_embeds[:, toks.argmax()] = prompt_embeds[:, toks.argmax()] + self.avg_diff[0] * scale
-                else:
-                    normed_prompt_embeds = prompt_embeds / prompt_embeds.norm(dim=-1, keepdim=True)
-                    sims = normed_prompt_embeds[0] @ normed_prompt_embeds[0].T
-                    if i == 0:
-                        weights = sims[toks.argmax(), :][None, :, None].repeat(1, 1, 768)
-
-                        standard_weights = torch.ones_like(weights)
-
-                        weights = standard_weights + (weights - standard_weights) * correlation_weight_factor
-                        prompt_embeds = prompt_embeds + (weights * self.avg_diff[0][None, :].repeat(1, self.pipe.tokenizer.model_max_length, 1) * scale)
-                    else:
-                        weights = sims[toks.argmax(), :][None, :, None].repeat(1, 1, 1280)
-
-                        standard_weights = torch.ones_like(weights)
-
-                        weights = standard_weights + (weights - standard_weights) * correlation_weight_factor
-                        prompt_embeds = prompt_embeds + (weights * self.avg_diff[1][None, :].repeat(1, self.pipe.tokenizer_2.model_max_length, 1) * scale)
-
-                bs_embed, seq_len, _ = prompt_embeds.shape
-                prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
-                if i < 2:
-                    clip_prompt_embeds_list.append(prompt_embeds)
-
-            clip_prompt_embeds = torch.concat(clip_prompt_embeds_list, dim=-1)
-            clip_pooled_prompt_embeds = torch.concat(clip_pooled_prompt_embeds_list, dim=-1)
-
-            clip_prompt_embeds = torch.nn.functional.pad(
-                clip_prompt_embeds, (0, t5_prompt_embed_shape - clip_prompt_embeds.shape[-1])
-            )
-
-            prompt_embeds = torch.cat([clip_prompt_embeds, prompt_embeds], dim=-2)
-
-
-
-            torch.manual_seed(seed)
-            image = self.pipe(prompt_embeds=prompt_embeds, pooled_prompt_embeds=clip_pooled_prompt_embeds,
-                         **pipeline_kwargs).images[0]
-
-        return image