Update ledits/pipeline_leditspp_stable_diffusion_xl.py

ledits/pipeline_leditspp_stable_diffusion_xl.py

@@ -613,11 +613,10 @@ class LEditsPPPipelineStableDiffusionXL(
                 else:
                     # "2" because SDXL always indexes from the penultimate layer.
                     edit_concepts_embeds = edit_concepts_embeds.hidden_states[-(clip_skip + 2)]
-                
-                print("SHALOM???")
-                if avg_diff is not None and avg_diff_2 is not None:
-                    #scale=3
-                    print("SHALOM")
+
+
+                if avg_diff is not None:
+
                     normed_prompt_embeds = edit_concepts_embeds / edit_concepts_embeds.norm(dim=-1, keepdim=True)
                     sims = normed_prompt_embeds[0] @ normed_prompt_embeds[0].T
                     if i == 0:
@@ -626,14 +625,26 @@ class LEditsPPPipelineStableDiffusionXL(
                         standard_weights = torch.ones_like(weights)
 
                         weights = standard_weights + (weights - standard_weights) * correlation_weight_factor
-                        edit_concepts_embeds = edit_concepts_embeds + (weights * avg_diff[None, :].repeat(1,tokenizer.model_max_length, 1) * scale)
+                        edit_concepts_embeds = edit_concepts_embeds + (
+                                    weights * avg_diff[0][None, :].repeat(1, tokenizer.model_max_length, 1) * scale)
+
+                        if avg_diff_2nd is not None:
+                            edit_concepts_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
-                        edit_concepts_embeds = edit_concepts_embeds + (weights * avg_diff_2[None, :].repeat(1, tokenizer.model_max_length, 1) * scale)
+                        edit_concepts_embeds = edit_concepts_embeds + (
+                                    weights * avg_diff[1][None, :].repeat(1, tokenizer.model_max_length, 1) * scale)
+                        if avg_diff_2nd is not None:
+                            edit_concepts_embeds += (weights * avg_diff_2nd[1][None, :].repeat(1,
+                                                                                            self.pipe.tokenizer_2.model_max_length,
+                                                                                            1) * scale_2nd)
+
 
                 edit_prompt_embeds_list.append(edit_concepts_embeds)
                 i+=1