anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Hardware][MUSA] Make MiniMax H3 conditioned VAE RNG device-aware (#5703)

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR authored and GitHub committed 2026-08-05 03:11:46 -07:00
1 parent 7fb62db5c2
commit 2eada79b99
3 files changed
+156 -7
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/minimax_h3/vae.py
+7 -6
View File
@@ -117,6 +117,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
self.use_tiling = True
self.use_slicing = False
self.parallel_size = 1
self.device_module = torch.get_device_module()
def set_parallel_size(
self,
@@ -166,13 +167,13 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
previous_dtype = parameter.dtype
if previous_dtype != torch.float32:
self.to(torch.float32)
devices = [parameter.device] if parameter.device.type == "cuda" else []
devices = [parameter.device] if parameter.device.type != "cpu" else []
try:
with torch.random.fork_rng(devices=devices):
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
for device in devices:
with torch.cuda.device(device):
torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
with self.device_module.device(device):
self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
latent = self.model.encode_images(
image,
use_fp16_latent=True,
@@ -210,13 +211,13 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
previous_dtype = parameter.dtype
if previous_dtype != torch.float32:
self.to(torch.float32)
devices = [parameter.device] if parameter.device.type == "cuda" else []
devices = [parameter.device] if parameter.device.type != "cpu" else []
try:
with torch.random.fork_rng(devices=devices):
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
for device in devices:
with torch.cuda.device(device):
torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
with self.device_module.device(device):
self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
latent = self.model.encode_videos(
frames,
use_fp16_latent=True,