Gsq/dev train

[gushiqiao]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces support for the Wan2.2 TI2V 5B model, adding new training and inference configurations, model definitions, and VAE components. It implements sequence parallel (SP) support across the attention layers and dataset loaders, introduces new DMD video trainers (VideoDmdTrainer and VideoArDmdTrainer), and adds a utility script to pre-compute teacher-forcing caches. A critical issue was identified in the VAE implementation where comparing a PyTorch tensor directly with a string placeholder could raise a RuntimeError due to ambiguous boolean evaluation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Comparing a PyTorch tensor directly with a string using == or != returns a tensor of booleans. When evaluated in an if condition, this will raise a RuntimeError: Boolean value of Tensor with more than one value is ambiguous. Since feat_cache[idx] can be either "Rep" (a string) or a PyTorch tensor, we should use isinstance(feat_cache[idx], str) to safely check if it is the placeholder string before performing any comparisons.

Suggested change

	if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
	# cache last frame of last two chunk
	cache_x = torch.cat(
	[
	feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
	cache_x,
	],
	dim=2,
	)
	if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
	cache_x = torch.cat(
	[torch.zeros_like(cache_x).to(cache_x.device), cache_x],
	dim=2,
	)
	if feat_cache[idx] == "Rep":
	x = self.time_conv(x)
	else:
	x = self.time_conv(x, feat_cache[idx])
	is_rep = isinstance(feat_cache[idx], str) and feat_cache[idx] == "Rep"
	if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not is_rep:
	# cache last frame of last two chunk
	cache_x = torch.cat(
	[
	feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
	cache_x,
	],
	dim=2,
	)
	if cache_x.shape[2] < 2 and feat_cache[idx] is not None and is_rep:
	cache_x = torch.cat(
	[torch.zeros_like(cache_x).to(cache_x.device), cache_x],
	dim=2,
	)
	if is_rep:
	x = self.time_conv(x)
	else:
	x = self.time_conv(x, feat_cache[idx])