diff --git a/helios/diffusers_version/pipeline_helios_diffusers.py b/helios/diffusers_version/pipeline_helios_diffusers.py index 99e2dbd..b2a3ffc 100644 --- a/helios/diffusers_version/pipeline_helios_diffusers.py +++ b/helios/diffusers_version/pipeline_helios_diffusers.py @@ -21,6 +21,7 @@ import numpy as np import regex as re import torch import torch.nn.functional as F +from accelerate.utils import broadcast from transformers import AutoTokenizer, UMT5EncoderModel from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback @@ -697,6 +698,7 @@ class HeliosPipeline(DiffusionPipeline, HeliosLoraLoaderMixin): batch_size, channel, num_frames, height, width = latents.shape noise = self.sample_block_noise(batch_size, channel, num_frames, height, width, patch_size, device) noise = noise.to(device=device, dtype=transformer_dtype) + noise = broadcast(noise, from_process=0) latents = alpha * latents + beta * noise # To fix the block artifact if self.config.is_distilled: diff --git a/helios/pipelines/pipeline_helios.py b/helios/pipelines/pipeline_helios.py index 9df8022..aac0f37 100644 --- a/helios/pipelines/pipeline_helios.py +++ b/helios/pipelines/pipeline_helios.py @@ -21,6 +21,7 @@ from typing import Any, Callable, Dict, List, Optional, Union import regex as re import torch import torch.nn.functional as F +from accelerate.utils import broadcast from einops import rearrange from transformers import AutoTokenizer, UMT5EncoderModel @@ -670,6 +671,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): batch_size, channel, num_frames, height, width = latents.shape noise = self.sample_block_noise(batch_size, channel, num_frames, height, width) noise = noise.to(device=device, dtype=transformer_dtype) + noise = broadcast(noise, from_process=0) latents = alpha * latents + beta * noise # To fix the block artifact if use_dmd: diff --git a/helios/pipelines/pipeline_helios_ode.py b/helios/pipelines/pipeline_helios_ode.py index 3437517..de11191 100644 --- a/helios/pipelines/pipeline_helios_ode.py +++ b/helios/pipelines/pipeline_helios_ode.py @@ -21,6 +21,7 @@ from typing import Any, Callable, Dict, List, Optional, Union import regex as re import torch import torch.nn.functional as F +from accelerate.utils import broadcast from einops import rearrange from transformers import AutoTokenizer, UMT5EncoderModel @@ -677,6 +678,7 @@ class HeliosPipeline(DiffusionPipeline, WanLoraLoaderMixin): batch_size, channel, num_frames, height, width = latents.shape noise = self.sample_block_noise(batch_size, channel, num_frames, height, width) noise = noise.to(device=device, dtype=transformer_dtype) + noise = broadcast(noise, from_process=0) latents = alpha * latents + beta * noise # To fix the block artifact if use_dmd: