1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
This commit is contained in:
yiyixuxu
2024-10-26 05:27:57 +02:00
parent 303b47cc53
commit 1069d210e1
2 changed files with 9 additions and 8 deletions

View File

@@ -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]

View File

@@ -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