From da848b52dec333b4ec6c897bfdd31b303027d795 Mon Sep 17 00:00:00 2001 From: SHYuanBest Date: Sun, 8 Mar 2026 08:56:22 +0000 Subject: [PATCH] make style and quality --- .../pipeline_helios_diffusers.py | 11 +++++++--- helios/pipelines/pipeline_helios.py | 20 ++++++++++++++++--- helios/pipelines/pipeline_helios_ode.py | 20 ++++++++++++++++--- 3 files changed, 42 insertions(+), 9 deletions(-) diff --git a/helios/diffusers_version/pipeline_helios_diffusers.py b/helios/diffusers_version/pipeline_helios_diffusers.py index 4e322f9..3964662 100644 --- a/helios/diffusers_version/pipeline_helios_diffusers.py +++ b/helios/diffusers_version/pipeline_helios_diffusers.py @@ -491,9 +491,12 @@ class HeliosPipeline(DiffusionPipeline, HeliosLoraLoaderMixin): _, ph, pw = patch_size block_size = ph * pw - cov = torch.eye(block_size, device=device) * (1 + gamma) - torch.ones(block_size, block_size, device=device) * gamma + cov = ( + torch.eye(block_size, device=device) * (1 + gamma) + - torch.ones(block_size, block_size, device=device) * gamma + ) cov += torch.eye(block_size, device=device) * 1e-6 - cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. + cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. L = torch.linalg.cholesky(cov) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) @@ -703,7 +706,9 @@ class HeliosPipeline(DiffusionPipeline, HeliosLoraLoaderMixin): beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) batch_size, channel, num_frames, height, width = latents.shape - noise = self.sample_block_noise(batch_size, channel, num_frames, height, width, patch_size, device, generator) + noise = self.sample_block_noise( + batch_size, channel, num_frames, height, width, patch_size, device, generator + ) noise = noise.to(device=device, dtype=transformer_dtype) latents = alpha * latents + beta * noise # To fix the block artifact diff --git a/helios/pipelines/pipeline_helios.py b/helios/pipelines/pipeline_helios.py index 4575052..0a59cce 100644 --- a/helios/pipelines/pipeline_helios.py +++ b/helios/pipelines/pipeline_helios.py @@ -461,9 +461,12 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): _, ph, pw = patch_size block_size = ph * pw - cov = torch.eye(block_size, device=device) * (1 + gamma) - torch.ones(block_size, block_size, device=device) * gamma + cov = ( + torch.eye(block_size, device=device) * (1 + gamma) + - torch.ones(block_size, block_size, device=device) * gamma + ) cov += torch.eye(block_size, device=device) * 1e-6 - cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. + cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. L = torch.linalg.cholesky(cov) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) @@ -631,6 +634,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): transformer_dtype: torch.dtype = None, scheduler_type: str = "unipc", # unipc, euler use_dynamic_shifting: bool = False, + generator: torch.Generator | list[torch.Generator] | None = None, # ------------ CFG Zero ------------ use_cfg_zero_star: Optional[bool] = False, use_zero_init: Optional[bool] = True, @@ -692,7 +696,16 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) batch_size, channel, num_frames, height, width = latents.shape - noise = self.sample_block_noise(batch_size, channel, num_frames, height, width, self.transformer.config.patch_size, device, generator) + noise = self.sample_block_noise( + batch_size, + channel, + num_frames, + height, + width, + self.transformer.config.patch_size, + device, + generator, + ) noise = noise.to(device=device, dtype=transformer_dtype) latents = alpha * latents + beta * noise # To fix the block artifact @@ -1362,6 +1375,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): transformer_dtype=transformer_dtype, scheduler_type=scheduler_type, use_dynamic_shifting=use_dynamic_shifting, + generator=generator, # ------------ CFG Zero ------------ use_cfg_zero_star=use_cfg_zero_star, use_zero_init=use_zero_init, diff --git a/helios/pipelines/pipeline_helios_ode.py b/helios/pipelines/pipeline_helios_ode.py index 5bd7a3b..e4c7329 100644 --- a/helios/pipelines/pipeline_helios_ode.py +++ b/helios/pipelines/pipeline_helios_ode.py @@ -460,9 +460,12 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): _, ph, pw = patch_size block_size = ph * pw - cov = torch.eye(block_size, device=device) * (1 + gamma) - torch.ones(block_size, block_size, device=device) * gamma + cov = ( + torch.eye(block_size, device=device) * (1 + gamma) + - torch.ones(block_size, block_size, device=device) * gamma + ) cov += torch.eye(block_size, device=device) * 1e-6 - cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. + cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. L = torch.linalg.cholesky(cov) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) @@ -630,6 +633,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): transformer_dtype: torch.dtype = None, scheduler_type: str = "unipc", # unipc, euler use_dynamic_shifting: bool = False, + generator: torch.Generator | list[torch.Generator] | None = None, # ------------ CFG Zero ------------ use_cfg_zero_star: Optional[bool] = False, use_zero_init: Optional[bool] = True, @@ -699,7 +703,16 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) batch_size, channel, num_frames, height, width = latents.shape - noise = self.sample_block_noise(batch_size, channel, num_frames, height, width, self.transformer.config.patch_size, device, generator) + noise = self.sample_block_noise( + batch_size, + channel, + num_frames, + height, + width, + self.transformer.config.patch_size, + device, + generator, + ) noise = noise.to(device=device, dtype=transformer_dtype) latents = alpha * latents + beta * noise # To fix the block artifact @@ -1417,6 +1430,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): transformer_dtype=transformer_dtype, scheduler_type=scheduler_type, use_dynamic_shifting=use_dynamic_shifting, + generator=generator, # ------------ CFG Zero ------------ use_cfg_zero_star=use_cfg_zero_star, use_zero_init=use_zero_init,