From 1069d210e1b9e84a366cdc7a13965626ea258178 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 26 Oct 2024 05:27:57 +0200 Subject: [PATCH] test --- src/diffusers/pipelines/mochi/pipeline_mochi.py | 3 +-- .../scheduling_flow_match_euler_discrete.py | 14 ++++++++------ 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/diffusers/pipelines/mochi/pipeline_mochi.py b/src/diffusers/pipelines/mochi/pipeline_mochi.py index b5bb3e3bbd..c0c99a7274 100644 --- a/src/diffusers/pipelines/mochi/pipeline_mochi.py +++ b/src/diffusers/pipelines/mochi/pipeline_mochi.py @@ -660,11 +660,10 @@ class MochiPipeline(DiffusionPipeline): # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - print(t) noise_pred = self.transformer( hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, - timestep=1000 - timestep, + timestep=timestep, encoder_attention_mask=prompt_attention_mask, return_dict=False, )[0] diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index c869c4f5fc..c1096dbe0c 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -205,9 +205,15 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) timesteps = sigmas * self.config.num_train_timesteps - self.timesteps = timesteps.to(device=device) - self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + if self.config.invert_sigmas: + sigmas = 1.0 - sigmas + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + self.timesteps = timesteps.to(device=device) + self.sigmas = sigmas self._step_index = None self._begin_index = None @@ -295,10 +301,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): sigma = self.sigmas[self.step_index] sigma_next = self.sigmas[self.step_index + 1] - if self.config.invert_sigmas: - print("inverting") - sigma, sigma_next = sigma_next, sigma - prev_sample = sample + (sigma_next - sigma) * model_output # Cast sample back to model compatible dtype