Update src/pipeline_pe_clone.py

src/pipeline_pe_clone.py

@@ -48,14 +48,24 @@ def prepare_latent_image_ids_2(height, width, device, dtype):
     latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width//2, device=device)[None, :]   # x坐标
     return latent_image_ids
 
+# def position_encoding_clone(batch_size, original_height, original_width, device, dtype):
+#     latent_image_ids = prepare_latent_image_ids_2(original_height, original_width, device, dtype)
+#     latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+#     latent_image_ids = latent_image_ids.reshape(
+#             latent_image_id_height * latent_image_id_width, latent_image_id_channels
+#         )
+#     cond_latent_image_ids = latent_image_ids
+#     latent_image_ids = torch.concat([latent_image_ids, cond_latent_image_ids], dim=-2)
+#     return latent_image_ids
+
 def position_encoding_clone(batch_size, original_height, original_width, device, dtype):
     latent_image_ids = prepare_latent_image_ids_2(original_height, original_width, device, dtype)
-    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
-    latent_image_ids = latent_image_ids.reshape(
-            latent_image_id_height * latent_image_id_width, latent_image_id_channels
-        )
-    cond_latent_image_ids = latent_image_ids
-    latent_image_ids = torch.concat([latent_image_ids, cond_latent_image_ids], dim=-2)
+    # latent_image_ids shape: (H//2, W//2, 3)
+    latent_image_ids = latent_image_ids.reshape(-1, latent_image_ids.shape[-1])  # shape: (N, 3)
+    # Duplicate by concatenating along the token dimension (if that’s intended)
+    latent_image_ids = torch.cat([latent_image_ids, latent_image_ids], dim=0)  # shape: (2*N, 3)
+    # Add a batch dimension and repeat it for each sample in the batch
+    latent_image_ids = latent_image_ids.unsqueeze(0).repeat(batch_size, 1, 1)  # shape: (B, 2*N, 3)
     return latent_image_ids
 
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents