[Feature] Add Cosmos2 i2v pipeline #837
Conversation
| _class_name: str = "AutoencoderKLWan" | ||
| _diffusers_version: str = "0.34.0.dev0" | ||
| _name_or_path: str = "" |
There was a problem hiding this comment.
These fields should be pop/removed in the loader so can be removed. You can refer to how wan's vae config is defined
|
|
||
| def __post_init__(self): | ||
| self.blend_num_frames = (self.tile_sample_min_num_frames - | ||
| self.tile_sample_stride_num_frames) * 2 No newline at end of file |
|
|
||
|
|
||
| @dataclass | ||
| class CosmosVideoConfigFixed(CosmosVideoConfig): |
There was a problem hiding this comment.
Why is this needed?
| class CosmosVideoConfigFixed(CosmosVideoConfig): | ||
| """Fixed Cosmos Video Config that matches original Cosmos2 Video2World configuration.""" | ||
|
|
||
| def update_model_arch(self, config: dict) -> None: |
There was a problem hiding this comment.
Did you align against diffusers or original cosmos?
fastvideo/configs/sample/cosmos.py
Outdated
|
|
||
|
|
||
| @dataclass | ||
| class CosmosTeaCacheParams(CacheParams): |
There was a problem hiding this comment.
This can be removed for now, but you should add a default for cosmos sampling params. Refer to wan
fastvideo/layers/layernorm.py
Outdated
| if self.has_weight: | ||
| self.weight = nn.Parameter(self.weight) | ||
|
|
||
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
can you change this to forward_diffusers?
kevin314
force-pushed
the
klin/cosmos2
branch
from
ca83c7e to
8b89aff
Compare
| return x.flatten(-2) | ||
|
|
||
|
|
||
| def apply_rotary_emb( |
There was a problem hiding this comment.
Maybe we should consider making this cosmos2 specific?
| return imgs | ||
|
|
||
|
|
||
| def get_timestep_embedding( |
There was a problem hiding this comment.
Maybe also make the architecture specific
| bs, seq_len, _ = hidden_states.shape | ||
| num_seqs = bs | ||
| n, c = self.n_heads, self.d_model // self.total_num_heads | ||
| #n, c = self.n_heads, self.d_model // self.total_num_heads |
| qkv, _ = self.qkv_proj(hidden_states) | ||
| # Projection of 'own' hidden state (self-attention). No GQA here. | ||
| q, k, v = qkv.split(self.inner_dim, dim=-1) | ||
| #q, k, v = qkv.split(self.inner_dim, dim=-1) |
| timesteps_array = np.linspace(self._sigma_to_t(self.sigma_max), | ||
| self._sigma_to_t(self.sigma_min), | ||
| num_inference_steps) | ||
| t_max = self._sigma_to_t(self.sigma_max) |
There was a problem hiding this comment.
maybe make a separate version of this scheduler
| # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... | ||
| batch.generator = [ | ||
| torch.Generator("cpu").manual_seed(seed) for seed in seeds | ||
| torch.Generator(device="cpu").manual_seed(seed) for seed in seeds |
There was a problem hiding this comment.
maybe also separate these
There was a problem hiding this comment.
make arch specific
| self.world_size = self.fastvideo_args.num_gpus | ||
| self.shutting_down = False | ||
|
|
||
| # Initialize CUDA before setting up multiprocessing to ensure |
kevin314
force-pushed
the
klin/cosmos2
branch
from
275f0b2 to
1b7bfe5
Compare
2 participants
No description provided.