Wan sf platform

[Watebear]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces tensor parallel (TP) support for Flux2 models, sequence parallel (SP) support for LongCatImage models, platform-specific attention fallbacks for Hunyuan Video, and various platform-specific optimizations. The code review identifies several critical issues that need to be addressed: a potential shape mismatch in _npu_attention, a missing initialization of self.seq_p_group in LongCatImageTransformerModel leading to an AttributeError, a potential crash in WorldMirror attention when flash attention is unavailable, multi-node compatibility issues where global ranks are used as device ordinals instead of local ranks, and a missing explicit import of dist in the Hunyuan Video runner.

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]

In _npu_attention, out is initialized with shape (batch * seqlen, heads * dim). If out_unpad returned by the NPU flash attention has shape [total_unpadded_tokens, heads, dim], assigning it to out[indices] will raise a shape mismatch error. To make this robust and prevent runtime crashes, initialize out with shape (batch * seqlen, heads, dim) and reshape out_unpad to (-1, heads, dim) before assignment.

Suggested change

	out = torch.zeros(batch * seqlen, heads * dim, dtype=out_unpad.dtype, device=out_unpad.device)
	out[indices] = out_unpad
	return out.reshape(batch, seqlen, heads, dim)
	out = torch.zeros(batch * seqlen, heads, dim, dtype=out_unpad.dtype, device=out_unpad.device)
	out[indices] = out_unpad.reshape(-1, heads, dim)
	return out.reshape(batch, seqlen, heads, dim)

[gemini-code-assist]

self.seq_p_group is used in _seq_parallel_pre_process and _seq_parallel_post_process but is never initialized in LongCatImageTransformerModel.__init__. This will raise an AttributeError when seq_parallel is enabled. Initialize self.seq_p_group in __init__ when seq_parallel is True.

        if self.config.get("seq_parallel", False):
            self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p")
        else:
            self.seq_p_group = None
        self._init_infer_class()
        self._init_weights()
        self._init_infer()

[gemini-code-assist]

Since flash_attn_func_v2 can now be None when flash_attn is not installed, the attention forward pass in _apply_attention (lines 60-84) will crash with a TypeError: 'NoneType' object is not callable when running in bfloat16 or float16 precision. Please update _apply_attention to check if flash attention is actually available (e.g., _USE_FLASH_ATTN_V3 or flash_attn_func_v2 is not None) before attempting to use it, and otherwise fallback to F.scaled_dot_product_attention.

[gemini-code-assist]

Using dist.get_rank() directly as the device ordinal in torch.device(f"{ai_device}:{dist.get_rank()}") will fail on multi-node setups where the global rank exceeds the number of GPUs per node (e.g., global rank 8 on a node with 8 GPUs). Use the local rank (e.g., from the LOCAL_RANK environment variable or via dist.get_rank() % device_count) to ensure multi-node compatibility.

Suggested change

	if dist.is_initialized():
	return torch.device(f"{ai_device}:{dist.get_rank()}")
	return torch.device(ai_device)
	if dist.is_initialized():
	import os
	local_rank = int(os.environ.get("LOCAL_RANK", dist.get_rank() % torch.cuda.device_count()))
	return torch.device(f"{ai_device}:{local_rank}")
	return torch.device(ai_device)

[gemini-code-assist]

Using dist.get_rank() directly as the device ordinal in torch.device(f"{AI_DEVICE}:{dist.get_rank()}") will fail on multi-node setups where the global rank exceeds the number of GPUs per node. Use the local rank (e.g., from the LOCAL_RANK environment variable or via dist.get_rank() % device_count) to ensure multi-node compatibility.

Suggested change

	if dist.is_initialized():
	return torch.device(f"{AI_DEVICE}:{dist.get_rank()}")
	return torch.device(AI_DEVICE)
	if dist.is_initialized():
	import os
	local_rank = int(os.environ.get("LOCAL_RANK", dist.get_rank() % torch.cuda.device_count()))
	return torch.device(f"{AI_DEVICE}:{local_rank}")
	return torch.device(AI_DEVICE)

[gemini-code-assist]

The module dist (from torch.distributed) is used in this file but is not explicitly imported at the top of the file. Relying on wildcard imports (like from lightx2v.utils.utils import *) for standard/framework libraries is fragile and can lead to NameErrors if the utility module changes. Please explicitly import torch.distributed as dist at the top of the file.