Update src/pipeline_pe_clone.py

src/pipeline_pe_clone.py

@@ -48,24 +48,14 @@ 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_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)
+    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
 
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents