make style and quality

This commit is contained in:
SHYuanBest
2026-03-08 08:56:22 +00:00
parent 46a0c8874a
commit da848b52de
3 changed files with 42 additions and 9 deletions
@@ -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
+17 -3
View File
@@ -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,
+17 -3
View File
@@ -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,