Add Cosmos3 Edge model support#47181
Conversation
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Some initial comments from my side. Imo my biggest point is to refactor the vision tower as it feels a bit forced into the siglip2 model but it leads to very weird patterns. Ideally we should have a ViT like model structure that handles variable sequence lengths (by also preparing the vision inputs accordingly).
Other than that, imo the text model could be a way more simple Llama like
| torch.as_tensor(rope_deltas, dtype=torch.long, device=input_ids.device).unsqueeze(1), | ||
| ) | ||
|
|
||
| def compute_3d_position_ids( |
There was a problem hiding this comment.
Same here then to pass mm token type ids along
We can simplify by that
| position_ids = position_ids + self.rope_deltas.to(device=inputs_embeds.device).unsqueeze(0) | ||
| return position_ids | ||
|
|
||
| @can_return_tuple |
|
|
||
|
|
||
| @auto_docstring | ||
| class Cosmos3EdgeImageProcessor(TorchvisionBackend): |
There was a problem hiding this comment.
Could be similar to qwen vl? I think the only core difference is the reshapes happening in the model atm but we should refactor the vision model either way
There was a problem hiding this comment.
Let's move the processors to modular imo. These are very similar to qwen2 vl. The smart resize also looks identical to me (just slightly different refactor)
For functions you can just import it in modular and it resolves it by itself
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Structurally much better now imo, I think we can now start aligning to existing VLMs a bit more and use a tad more modular 🤗
| def apply_interleaved_mrope(freqs, mrope_section): | ||
| freqs_t = freqs[0] | ||
| for dim, offset in enumerate((1, 2), start=1): | ||
| length = mrope_section[dim] * 3 | ||
| freqs_t[..., slice(offset, length, 3)] = freqs[dim, ..., slice(offset, length, 3)] | ||
| return freqs_t |
There was a problem hiding this comment.
Rebump, not super important but still
|
|
||
|
|
||
| @auto_docstring | ||
| class Cosmos3EdgeImageProcessor(TorchvisionBackend): |
There was a problem hiding this comment.
Let's move the processors to modular imo. These are very similar to qwen2 vl. The smart resize also looks identical to me (just slightly different refactor)
For functions you can just import it in modular and it resolves it by itself
| if per_image_kwargs is not None and image_index < len(per_image_kwargs): | ||
| image_kwargs = per_image_kwargs[image_index] or {} |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, cosmos3_edge |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Looking better, it's getting to the smaller details. One bigger point I only noticed now because I didnt check the tests too closely before: We need to use the respective tester/test classes; currently we do not which means we skip a lot of our mixin tests
| # The checkpoint uses a learned square reference grid, interpolated independently for every packed frame. | ||
| target_dtype = self.patch_embedding.weight.dtype | ||
| patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) | ||
| positional_embeddings = self.position_embedding.weight.reshape( | ||
| self.position_embedding_size, self.position_embedding_size, -1 | ||
| ) | ||
| positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0) | ||
| source_dtype = positional_embeddings.dtype | ||
| if positional_embeddings.device.type == "cpu": | ||
| positional_embeddings = positional_embeddings.float() | ||
|
|
||
| position_chunks = [] | ||
| merge_size = self.config.spatial_merge_size | ||
| for temporal, height, width in grid_thw.tolist(): | ||
| resized_embeddings = F.interpolate( | ||
| positional_embeddings, | ||
| size=(height, width), | ||
| mode="bilinear", | ||
| align_corners=False, | ||
| antialias=True, | ||
| ) | ||
| resized_embeddings = resized_embeddings.squeeze(0).permute(1, 2, 0).to(source_dtype) | ||
| resized_embeddings = resized_embeddings.reshape( | ||
| height // merge_size, | ||
| merge_size, | ||
| width // merge_size, | ||
| merge_size, | ||
| -1, | ||
| ) | ||
| # Preserve the processor's block-major 2x2 patch order before the projector groups adjacent patches. | ||
| resized_embeddings = resized_embeddings.transpose(1, 2).reshape(height * width, -1) | ||
| position_chunks.append(resized_embeddings.repeat(temporal, 1)) | ||
|
|
||
| position_embeddings = torch.cat(position_chunks, dim=0) | ||
| torch_compilable_check( | ||
| patch_embeds.shape[0] == position_embeddings.shape[0], | ||
| "The packed visual patch count does not match `grid_thw`.", | ||
| ) | ||
| return patch_embeds + position_embeddings |
There was a problem hiding this comment.
Ok I noticed that we no longer really follow the functions exactly shouldnt the resize function be overriden for the portion here that is used?
So the same split as in the original with forward -> resize_positional_embeddings
There was a problem hiding this comment.
Please check out our VLMModelTester and VLMModelTest, e.g.
Currently a lot of our mixins are otherwise skipped. They also already handle the config tests
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
atharvajoshi10
force-pushed
the
cosmos-3-edge-integration
branch
from
163d593 to
9c62eaf
Compare
2 participants
What does this PR do?
This PR adds a support for Cosmos3 Reasoner model (not released yet).