Update pipeline_kiwi_edit.py

pipeline_kiwi_edit.py

@@ -63,10 +63,9 @@ class KiwiEditPipeline(DiffusionPipeline):
         mllm_encoder: MLLMEncoder - Qwen2.5-VL MLLM with learnable queries.
         processor: AutoProcessor - Qwen2.5-VL processor/tokenizer bundle.
         source_embedder: ConditionalEmbedder - VAE source conditioning.
-        ref_embedder: ConditionalEmbedder - VAE reference conditioning.
     """
 
-    model_cpu_offload_seq = "mllm_encoder->source_embedder->ref_embedder->transformer->vae"
+    model_cpu_offload_seq = "mllm_encoder->source_embedder->transformer->vae"
 
     def __init__(
         self,
@@ -75,7 +74,6 @@ class KiwiEditPipeline(DiffusionPipeline):
         scheduler,
         mllm_encoder,
         source_embedder,
-        ref_embedder,
         processor=None,
     ):
         super().__init__()
@@ -89,7 +87,6 @@ class KiwiEditPipeline(DiffusionPipeline):
             mllm_encoder=mllm_encoder,
             processor=processor,
             source_embedder=source_embedder,
-            ref_embedder=ref_embedder,
         )
         if processor is not None:
             self.mllm_encoder.processor = processor
@@ -229,29 +226,6 @@ class KiwiEditPipeline(DiffusionPipeline):
         # --- 3D RoPE frequencies (real-valued cos/sin format) ---
         rotary_emb = _build_rope_3d(t.rope, f, h, w, device)
 
-        # --- Reference image conditioning ---
-        vae_ref_input_length = 0
-        if vae_ref_image is not None:
-            if len(vae_ref_image) > 1:
-                vae_ref = torch.cat(vae_ref_image, dim=2)  # concat along temporal
-            else:
-                vae_ref = vae_ref_image[0]
-
-            vae_ref = self.ref_embedder(vae_ref)
-            ref_f, ref_h, ref_w = vae_ref.shape[2:]
-            vae_ref = rearrange(vae_ref, "b c f h w -> b (f h w) c").contiguous()
-
-            # Recompute RoPE for extended sequence (main + ref tokens)
-            total_f = f + ref_f
-            rotary_emb = _build_rope_3d(t.rope, total_f, h, w, device)
-
-            vae_ref_input_length = vae_ref.shape[1]
-
-            if self.ref_embedder.config.ref_pad_first:
-                x = torch.cat([vae_ref, x], dim=1)
-            else:
-                x = torch.cat([x, vae_ref], dim=1)
-
         # --- Transformer blocks ---
         for block in t.blocks:
             x = block(x, context, t_mod, rotary_emb)
@@ -267,13 +241,6 @@ class KiwiEditPipeline(DiffusionPipeline):
         x = (t.norm_out(x.float()) * (1 + scale) + shift).type_as(x)
         x = t.proj_out(x)
 
-        # --- Remove ref tokens from output ---
-        if vae_ref_image is not None and vae_ref_input_length > 0:
-            if self.ref_embedder.config.ref_pad_first:
-                x = x[:, vae_ref_input_length:, :]
-            else:
-                x = x[:, :-vae_ref_input_length, :]
-
         # --- Unpatchify ---
         patch_size = t.config.patch_size
         x = rearrange(
@@ -410,15 +377,15 @@ class KiwiEditPipeline(DiffusionPipeline):
             vae_source_input = vae_source_input.to(dtype=dtype)
 
         # --- 7. Encode reference images ---
-        vae_ref_image = None
-        if ref_image is not None:
-            vae_ref_image = []
-            for item in ref_image:
-                target_size = (width, height)
-                item = ImageOps.pad(item, target_size, color="white", centering=(0.5, 0.5))
-                ref_tensor = self._preprocess_video([item], dtype=torch.float32, device=device)
-                ref_latent = self.vae.encode(ref_tensor).latent_dist.sample()
-                vae_ref_image.append(ref_latent.to(dtype=dtype))
+        # vae_ref_image = None
+        # if ref_image is not None:
+        #     vae_ref_image = []
+        #     for item in ref_image:
+        #         target_size = (width, height)
+        #         item = ImageOps.pad(item, target_size, color="white", centering=(0.5, 0.5))
+        #         ref_tensor = self._preprocess_video([item], dtype=torch.float32, device=device)
+        #         ref_latent = self.vae.encode(ref_tensor).latent_dist.sample()
+        #         vae_ref_image.append(ref_latent.to(dtype=dtype))
 
         # --- 8. Handle input_video (video-to-video) ---
         if input_video is not None:
@@ -439,7 +406,7 @@ class KiwiEditPipeline(DiffusionPipeline):
                 timestep=timestep,
                 context=context,
                 vae_source_input=vae_source_input,
-                vae_ref_image=vae_ref_image,
+                # vae_ref_image=vae_ref_image,
                 sigmas=sigmas,
                 timesteps_schedule=timesteps,
             )