(self,
x: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
clip_feature: Optional[torch.Tensor] = None,
y: Optional[torch.Tensor] = None,
use_gradient_checkpointing: bool = False,
use_gradient_checkpointing_offload: bool = False,
**kwargs,
)
| 40 | return x_out.to(x.dtype) |
| 41 | |
| 42 | def usp_dit_forward(self, |
| 43 | x: torch.Tensor, |
| 44 | timestep: torch.Tensor, |
| 45 | context: torch.Tensor, |
| 46 | clip_feature: Optional[torch.Tensor] = None, |
| 47 | y: Optional[torch.Tensor] = None, |
| 48 | use_gradient_checkpointing: bool = False, |
| 49 | use_gradient_checkpointing_offload: bool = False, |
| 50 | **kwargs, |
| 51 | ): |
| 52 | t = self.time_embedding( |
| 53 | sinusoidal_embedding_1d(self.freq_dim, timestep)) |
| 54 | t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) |
| 55 | context = self.text_embedding(context) |
| 56 | |
| 57 | if self.has_image_input: |
| 58 | x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) |
| 59 | clip_embdding = self.img_emb(clip_feature) |
| 60 | context = torch.cat([clip_embdding, context], dim=1) |
| 61 | |
| 62 | x, (f, h, w) = self.patchify(x) |
| 63 | |
| 64 | freqs = torch.cat([ |
| 65 | self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| 66 | self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| 67 | self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) |
| 68 | ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) |
| 69 | |
| 70 | def create_custom_forward(module): |
| 71 | def custom_forward(*inputs): |
| 72 | return module(*inputs) |
| 73 | return custom_forward |
| 74 | |
| 75 | # Context Parallel |
| 76 | x = torch.chunk( |
| 77 | x, get_sequence_parallel_world_size(), |
| 78 | dim=1)[get_sequence_parallel_rank()] |
| 79 | |
| 80 | for block in self.blocks: |
| 81 | if self.training and use_gradient_checkpointing: |
| 82 | if use_gradient_checkpointing_offload: |
| 83 | with torch.autograd.graph.save_on_cpu(): |
| 84 | x = torch.utils.checkpoint.checkpoint( |
| 85 | create_custom_forward(block), |
| 86 | x, context, t_mod, freqs, |
| 87 | use_reentrant=False, |
| 88 | ) |
| 89 | else: |
| 90 | x = torch.utils.checkpoint.checkpoint( |
| 91 | create_custom_forward(block), |
| 92 | x, context, t_mod, freqs, |
| 93 | use_reentrant=False, |
| 94 | ) |
| 95 | else: |
| 96 | x = block(x, context, t_mod, freqs) |
| 97 | |
| 98 | x = self.head(x, t) |
| 99 |
nothing calls this directly
no test coverage detected