Edit attention map selection upon num_edit_tokens

pipeline_semantic_stable_diffusion_img2img_solver.py

@@ -928,11 +928,11 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
                                 from_where=["up", "down"],
                                 is_cross=True,
                                 select=text_cross_attention_maps.index(editing_prompt[c]),
-                            )
-                            attn_map = out[:, :, :, 1:1 + num_edit_tokens[c]]  # 0 -> startoftext
+                            )                            
+                            attn_map = out[:, :, :, 1:1 + num_edit_tokens[self.batch_size*c]]  # 0 -> startoftext
 
                             # average over all tokens
-                            assert (attn_map.shape[3] == num_edit_tokens[c])
+                            assert (attn_map.shape[3] == num_edit_tokens[self.batch_size*c])
                             attn_map = torch.sum(attn_map, dim=3)
 
                             # gaussian_smoothing