Skip to content

Commit b7b1165

Browse files
helloyongyanghelloyongyang
committed
Update sheduler generator. If generator is not None, use existing generator for latents.
1 parent 181846b commit b7b1165

12 files changed

Lines changed: 73 additions & 17 deletions

File tree

‎lightx2v/models/schedulers/bagel/scheduler.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import torch
2+
from loguru import logger
23

34
from lightx2v.models.schedulers.scheduler import BaseScheduler
45

@@ -55,7 +56,10 @@ def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids,
5556

5657
query_curr = curr = 0
5758
seed = int(seed if seed is not None else self.config.get("seed", 42))
58-
self.generator = torch.Generator(device="cpu").manual_seed(seed)
59+
if self.generator is None:
60+
self.generator = torch.Generator(device="cpu").manual_seed(seed)
61+
else:
62+
logger.info(f"Generator is not None, using existing generator for latents")
5963

6064
for (H, W), curr_kvlen, curr_position_id in zip(image_sizes, curr_kvlens, curr_rope):
6165
packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
@@ -137,7 +141,10 @@ def prepare_vae_latent_cfg(self, curr_kvlens, curr_rope, image_sizes):
137141
return generation_input
138142

139143
def prepare(self):
140-
self.generator = torch.Generator().manual_seed(42)
144+
if self.generator is None:
145+
self.generator = torch.Generator().manual_seed(42)
146+
else:
147+
logger.info(f"Generator is not None, using existing generator for latents")
141148
self.set_timesteps()
142149
self.generation_input = None
143150
self.generation_input_cfg_text = None

‎lightx2v/models/schedulers/flux2/scheduler.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import numpy as np
66
import torch
7+
from loguru import logger
78

89
try:
910
from diffusers.pipelines.flux2.pipeline_flux2 import compute_empirical_mu
@@ -76,7 +77,10 @@ def __init__(self, config):
7677

7778
def prepare(self, input_info):
7879
self.input_info = input_info
79-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
80+
if self.generator is None:
81+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
82+
else:
83+
logger.info(f"Generator is not None, using existing generator for latents")
8084

8185
if hasattr(input_info, "latent_image_ids"):
8286
self.latent_image_ids = input_info.latent_image_ids

‎lightx2v/models/schedulers/hunyuan_video/scheduler.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import torch
22
import torch.distributed as dist
33
from einops import rearrange
4+
from loguru import logger
45
from torch.nn import functional as F
56

67
from lightx2v.models.networks.hunyuan_video.infer.posemb_layers import get_nd_rotary_pos_embed
@@ -50,7 +51,10 @@ def prepare(self, seed, latent_shape, image_encoder_output=None):
5051
self.cond_latents_concat, self.mask_concat = self._prepare_cond_latents_and_mask(self.config["task"], image_encoder_output["cond_latents"], self.latents, self.multitask_mask, self.reorg_token)
5152

5253
def prepare_latents(self, seed, latent_shape, dtype=torch.bfloat16):
53-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
54+
if self.generator is None:
55+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
56+
else:
57+
logger.info(f"Generator is not None, using existing generator for latents")
5458
self.latents = torch.randn(
5559
1,
5660
latent_shape[0],
@@ -172,7 +176,10 @@ def prepare(self, seed, latent_shape, lq_latents, upsampler, image_encoder_outpu
172176
self.zero_condition = zero_condition
173177

174178
def prepare_latents(self, seed, latent_shape, lq_latents, dtype=torch.bfloat16):
175-
self.generator = torch.Generator(device=lq_latents.device).manual_seed(seed)
179+
if self.generator is None:
180+
self.generator = torch.Generator(device=lq_latents.device).manual_seed(seed)
181+
else:
182+
logger.info(f"Generator is not None, using existing generator for latents")
176183
self.latents = torch.randn(
177184
1,
178185
latent_shape[0],

‎lightx2v/models/schedulers/longcat_image/scheduler.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import numpy as np
88
import torch
99
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
10+
from loguru import logger
1011
from torch import nn
1112

1213
from lightx2v.models.schedulers.scheduler import BaseScheduler
@@ -336,7 +337,10 @@ def set_timesteps(self):
336337

337338
def prepare(self, input_info):
338339
"""Prepare scheduler for inference."""
339-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
340+
if self.generator is None:
341+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
342+
else:
343+
logger.info(f"Generator is not None, using existing generator for latents")
340344
self.prepare_latents(input_info)
341345
self.set_timesteps()
342346

@@ -403,7 +407,10 @@ def prepare_i2i(self, input_info, input_image, vae):
403407
input_image: Input image tensor [B, C, H, W] (preprocessed)
404408
vae: VAE model for encoding
405409
"""
406-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
410+
if self.generator is None:
411+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
412+
else:
413+
logger.info(f"Generator is not None, using existing generator for latents")
407414
self.vae = vae
408415
self.prepare_latents(input_info)
409416
self.set_timesteps()

‎lightx2v/models/schedulers/ltx2/scheduler.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import einops
1616
import torch
17+
from loguru import logger
1718

1819
from lightx2v.models.schedulers.scheduler import BaseScheduler
1920
from lightx2v.utils.envs import *
@@ -415,7 +416,10 @@ def prepare(
415416
self._video_main_num_tokens = None
416417

417418
# Initialize generator
418-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
419+
if self.generator is None:
420+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
421+
else:
422+
logger.info(f"Generator is not None, using existing generator for latents")
419423

420424
# Prepare latents
421425
self.prepare_latents(

‎lightx2v/models/schedulers/motus/scheduler.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import numpy as np
22
import torch
3+
from loguru import logger
34

45
from lightx2v.models.schedulers.wan.scheduler import WanScheduler
56

@@ -45,7 +46,10 @@ def prepare(self, seed, latent_shape, image_encoder_output, action_shape):
4546
self.latents = self.video_latents.squeeze(0)
4647

4748
def prepare_latents(self, seed, latent_shape, dtype=torch.float32):
48-
self.generator = torch.Generator(device=self.config.get("device", self.vae_encoder_out.device)).manual_seed(seed)
49+
if self.generator is None:
50+
self.generator = torch.Generator(device=self.config.get("device", self.vae_encoder_out.device)).manual_seed(seed)
51+
else:
52+
logger.info(f"Generator is not None, using existing generator for latents")
4953
self.latents = torch.randn(
5054
latent_shape[0],
5155
latent_shape[1],

‎lightx2v/models/schedulers/qwen_image/scheduler.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import torch
1111
import torch.distributed as dist
1212
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
13+
from loguru import logger
1314
from torch import nn
1415
from torch.nn import functional as F
1516

@@ -572,10 +573,13 @@ def set_timesteps(self):
572573
self.num_warmup_steps = num_warmup_steps
573574

574575
def prepare(self, input_info):
575-
if self.config["task"] == "i2i":
576-
self.generator = torch.Generator().manual_seed(input_info.seed)
577-
elif self.config["task"] == "t2i":
578-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
576+
if self.generator is None:
577+
if self.config["task"] == "i2i":
578+
self.generator = torch.Generator().manual_seed(input_info.seed)
579+
elif self.config["task"] == "t2i":
580+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(input_info.seed)
581+
else:
582+
logger.info(f"Generator is not None, using existing generator for latents")
579583
self.prepare_latents(input_info)
580584
self.set_timesteps()
581585

‎lightx2v/models/schedulers/scheduler.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ def clear(self):
1111
class BaseScheduler:
1212
def __init__(self, config):
1313
self.config = config
14+
self.generator = None
1415
self.latents = None
1516
self.step_index = 0
1617
self.infer_steps = config["infer_steps"]

‎lightx2v/models/schedulers/wan/audio/scheduler.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ class EulerScheduler(WanScheduler):
1414
def __init__(self, config):
1515
self.config = config
1616
self.latents = None
17+
self.generator = None
1718
self.step_index = 0
1819
self.flag_df = False
1920
self.transformer_infer = None
@@ -76,7 +77,10 @@ def step_pre(self, step_index):
7677
)
7778

7879
def prepare_latents(self, seed, latent_shape, dtype=torch.float32):
79-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
80+
if self.generator is None:
81+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
82+
else:
83+
logger.info(f"Generator is not None, using existing generator for latents")
8084
self.latents = torch.randn(
8185
latent_shape[0],
8286
latent_shape[1],

‎lightx2v/models/schedulers/wan/changing_resolution/scheduler.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import torch
2+
from loguru import logger
23

34
from lightx2v_platform.base.global_var import AI_DEVICE
45

@@ -15,14 +16,18 @@ def __init__(self, config):
1516

1617
class WanScheduler4ChangingResolution:
1718
def __init__(self, config):
19+
self.generator = None
1820
if "resolution_rate" not in config:
1921
config["resolution_rate"] = [0.75]
2022
if "changing_resolution_steps" not in config:
2123
config["changing_resolution_steps"] = [config.infer_steps // 2]
2224
assert len(config["resolution_rate"]) == len(config["changing_resolution_steps"])
2325

2426
def prepare_latents(self, seed, latent_shape, dtype=torch.float32):
25-
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
27+
if self.generator is None:
28+
self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed)
29+
else:
30+
logger.info(f"Generator is not None, using existing generator for latents")
2631
self.latents_list = []
2732
for i in range(len(self.config["resolution_rate"])):
2833
self.latents_list.append(

0 commit comments

Comments
 (0)