Upload folder using huggingface_hub

app.py

@@ -13,6 +13,7 @@ from omegaconf import OmegaConf
 from PIL import Image
 from typing import List, Optional
 import tempfile
+from huggingface_hub import hf_hub_download
 import random
 
 from torchvision import transforms
@@ -176,7 +177,12 @@ def temporal_downsample_for_vlm(pil_frames: List[Image.Image], cfg) -> List[Imag
 print("Building LoomVideo model...")
 model = load_model(config)
 print("Loading LoomVideo checkpoint...")
-load_checkpoint(model, LOOMVIDEO_CKPT)
+ckpt_file = hf_hub_download(
+    repo_id=LOOMVIDEO_CKPT,
+    filename="Stage3/gen_model.pth",
+    repo_type="model",
+)
+load_checkpoint(model, ckpt_file)
 model.to(dtype=torch.bfloat16, device="cuda")
 model.eval()
 print("Model loaded successfully.")
@@ -205,6 +211,20 @@ def _get_guidance_scale_visual(guidance_scale_visual: Optional[float]) -> float:
     return OmegaConf.select(config, "generation.guidance_scale_visual", default=1.5)
 
 
+def _build_generate_kwargs(inputs, **extra):
+    """Build kwargs dict for model.generate() from processor outputs."""
+    kwargs = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+    }
+    # Pass through optional keys that the processor may produce
+    for key in ["pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw", "mm_token_type_ids"]:
+        if key in inputs:
+            kwargs[key] = inputs[key]
+    kwargs.update(extra)
+    return kwargs
+
+
 # ---------------------------------------------------------------------------
 # Inference functions
 # ---------------------------------------------------------------------------
@@ -261,18 +281,17 @@ def generate_t2v(
     else:
         generator = None
 
-    generate_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "negative_prompt": _get_negative_prompt(None),
-        "height": height,
-        "width": width,
-        "num_frames": num_frames,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": _get_guidance_scale(guidance_scale, is_edit=False),
-        "guidance_scale_visual": _get_guidance_scale_visual(guidance_scale_visual),
-        "generator": generator,
-    }
+    generate_kwargs = _build_generate_kwargs(
+        inputs,
+        negative_prompt=_get_negative_prompt(None),
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=False),
+        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
+        generator=generator,
+    )
 
     with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
         output = model.generate(**generate_kwargs)
@@ -348,23 +367,18 @@ def generate_edit(
     else:
         generator = None
 
-    generate_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "pixel_values": inputs.get("pixel_values"),
-        "pixel_values_videos": inputs.get("pixel_values_videos"),
-        "image_grid_thw": inputs.get("image_grid_thw"),
-        "video_grid_thw": inputs.get("video_grid_thw"),
-        "negative_prompt": _get_negative_prompt(None),
-        "height": height,
-        "width": width,
-        "num_frames": num_frames,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": _get_guidance_scale(guidance_scale, is_edit=True),
-        "guidance_scale_visual": _get_guidance_scale_visual(guidance_scale_visual),
-        "source_pixel_values": [source_tensor],
-        "generator": generator,
-    }
+    generate_kwargs = _build_generate_kwargs(
+        inputs,
+        negative_prompt=_get_negative_prompt(None),
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=True),
+        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
+        source_pixel_values=[source_tensor],
+        generator=generator,
+    )
 
     with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
         output = model.generate(**generate_kwargs)
@@ -433,24 +447,19 @@ def generate_ref_edit(
     else:
         generator = None
 
-    generate_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "pixel_values": inputs.get("pixel_values"),
-        "pixel_values_videos": inputs.get("pixel_values_videos"),
-        "image_grid_thw": inputs.get("image_grid_thw"),
-        "video_grid_thw": inputs.get("video_grid_thw"),
-        "negative_prompt": _get_negative_prompt(None),
-        "height": height,
-        "width": width,
-        "num_frames": num_frames,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": _get_guidance_scale(guidance_scale, is_edit=True),
-        "guidance_scale_visual": _get_guidance_scale_visual(guidance_scale_visual),
-        "source_pixel_values": [source_video_tensor],
-        "ref_pixel_values": ref_image_tensors,
-        "generator": generator,
-    }
+    generate_kwargs = _build_generate_kwargs(
+        inputs,
+        negative_prompt=_get_negative_prompt(None),
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=True),
+        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
+        source_pixel_values=[source_video_tensor],
+        ref_pixel_values=ref_image_tensors,
+        generator=generator,
+    )
 
     with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
         output = model.generate(**generate_kwargs)
@@ -524,21 +533,18 @@ def generate_mi2v(
     else:
         generator = None
 
-    generate_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "pixel_values": inputs.get("pixel_values"),
-        "image_grid_thw": inputs.get("image_grid_thw"),
-        "negative_prompt": _get_negative_prompt(None),
-        "height": height,
-        "width": width,
-        "num_frames": num_frames,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": _get_guidance_scale(guidance_scale, is_edit=False),
-        "guidance_scale_visual": _get_guidance_scale_visual(guidance_scale_visual),
-        "ref_pixel_values": ref_image_tensors,
-        "generator": generator,
-    }
+    generate_kwargs = _build_generate_kwargs(
+        inputs,
+        negative_prompt=_get_negative_prompt(None),
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=False),
+        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
+        ref_pixel_values=ref_image_tensors,
+        generator=generator,
+    )
 
     with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
         output = model.generate(**generate_kwargs)

src/models/transformers/loomvideo.py

@@ -846,6 +846,7 @@ class LoomVideo(ModelMixin):
         pixel_values_videos: Optional[torch.FloatTensor] = None,
         image_grid_thw: Optional[torch.LongTensor] = None,
         video_grid_thw: Optional[torch.LongTensor] = None,
+        mm_token_type_ids: Optional[torch.IntTensor] = None,
         negative_prompt: Optional[str] = None,
         height: int = 480,
         width: int = 832,
@@ -925,6 +926,7 @@ class LoomVideo(ModelMixin):
             pixel_values_videos=pixel_values_videos,
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
+            mm_token_type_ids=mm_token_type_ids,
         )
         all_hidden_states_full = outputs_full["hidden_states"]
 

src/models/transformers/qwen3vl.py

@@ -191,6 +191,7 @@ class Qwen3VLModel(Qwen3VLModel):
         image_grid_thw: Optional[torch.LongTensor] = None,
         video_grid_thw: Optional[torch.LongTensor] = None,
         cache_position: Optional[torch.LongTensor] = None,
+        mm_token_type_ids: Optional[torch.IntTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, Qwen3VLModelOutputWithPast]:
         """Forward pass that extracts visual features, applies DeepStack injection,
@@ -205,7 +206,12 @@ class Qwen3VLModel(Qwen3VLModel):
         video_mask = None
 
         if pixel_values is not None:
-            image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw)
+            _img_out = self.get_image_features(pixel_values, image_grid_thw)
+            if isinstance(_img_out, (list, tuple)):
+                image_embeds, deepstack_image_embeds = _img_out
+            else:
+                image_embeds = _img_out.pooler_output
+                deepstack_image_embeds = _img_out.deepstack_features
             image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
             image_mask, _ = self.get_placeholder_mask(
                 input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
@@ -213,7 +219,12 @@ class Qwen3VLModel(Qwen3VLModel):
             inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
 
         if pixel_values_videos is not None:
-            video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw)
+            _vid_out = self.get_video_features(pixel_values_videos, video_grid_thw)
+            if isinstance(_vid_out, (list, tuple)):
+                video_embeds, deepstack_video_embeds = _vid_out
+            else:
+                video_embeds = _vid_out.pooler_output
+                deepstack_video_embeds = _vid_out.deepstack_features
             video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
             _, video_mask = self.get_placeholder_mask(
                 input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
@@ -264,12 +275,24 @@ class Qwen3VLModel(Qwen3VLModel):
                 or (past_key_values is None or past_key_values.get_seq_length() == 0)
             )
             if (prefill_compiled_stage or prefill_noncompiled_stage) or self.rope_deltas is None:
-                position_ids, rope_deltas = self.get_rope_index(
-                    input_ids,
-                    image_grid_thw,
-                    video_grid_thw,
-                    attention_mask=attention_mask_tensor,
-                )
+                # Handle transformers 5.x which requires mm_token_type_ids
+                import inspect
+                rope_sig = inspect.signature(self.get_rope_index)
+                if 'mm_token_type_ids' in rope_sig.parameters:
+                    position_ids, rope_deltas = self.get_rope_index(
+                        input_ids,
+                        mm_token_type_ids if mm_token_type_ids is not None else torch.zeros_like(input_ids, dtype=torch.int32),
+                        image_grid_thw,
+                        video_grid_thw,
+                        attention_mask=attention_mask_tensor,
+                    )
+                else:
+                    position_ids, rope_deltas = self.get_rope_index(
+                        input_ids,
+                        image_grid_thw,
+                        video_grid_thw,
+                        attention_mask=attention_mask_tensor,
+                    )
                 self.rope_deltas = rope_deltas
             else:
                 batch_size, seq_length, _ = inputs_embeds.shape