Skip to content

Commit

Permalink
[Chore] fix vae annotation in mochi pipeline (#10585)
Browse files Browse the repository at this point in the history
fix vae annotation in mochi pipeline
  • Loading branch information
sayakpaul authored Jan 15, 2025
1 parent bba59fb commit c944f06
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions src/diffusers/pipelines/mochi/pipeline_mochi.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

from ...callbacks import MultiPipelineCallbacks, PipelineCallback
from ...loaders import Mochi1LoraLoaderMixin
from ...models.autoencoders import AutoencoderKL
from ...models.autoencoders import AutoencoderKLMochi
from ...models.transformers import MochiTransformer3DModel
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import (
Expand Down Expand Up @@ -151,8 +151,8 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
Conditional Transformer architecture to denoise the encoded video latents.
scheduler ([`FlowMatchEulerDiscreteScheduler`]):
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
vae ([`AutoencoderKLMochi`]):
Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
text_encoder ([`T5EncoderModel`]):
[T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
Expand All @@ -171,7 +171,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
def __init__(
self,
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
vae: AutoencoderKLMochi,
text_encoder: T5EncoderModel,
tokenizer: T5TokenizerFast,
transformer: MochiTransformer3DModel,
Expand Down

0 comments on commit c944f06

Please sign in to comment.