Fix Mochi Quality Issues

[DN6]

What does this PR do?

We're seeing some quality issues with Mochi due to missing upcasts and differences between how attention is handled in the original repo.

This PR:

1. Matches the transformer implementation 1:1 so that norms are upcast and run in the same precision as the original repo. 2. Changes the MochiAttnProcessor to match the original approach of dropping padding tokens.
2. Runs the CFG and Sampling step in FP32

I'll update the docs PR: #9934 with a guide on how to reproduce the original repo results exactly once this PR is merged.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

Mochi zeros out negative prompt embeds if they are an empty string.

[yiyixuxu]

thank you for the fix!!!!
so much work went into this PR:)

[yiyixuxu]

OK! But cc @a-r-r-o-w here. He has been following the mochi-fix PR and added the attention processors to model files
I guess we keep them here for now until we refactor and move them all together?

[a-r-r-o-w]

I think we should do the following going forward as design choice (just personal opinion so let's try to forumalate a plan for consistency):

• Apart from transformer model code, the transformer files will also contain all the relevant attention processor implementation. This helps with readability because you don't have to jump between files and because attention_processor.py is now > 5k lines
• If an alternate Attention class is required, let's keep it in the transformer file as well. These custom classes require some common methods that will probably not change between implementantions. For this, let's create a AttentionMixin class - for changing/getting attention processors, fusing, etc.
• If an attention processor is required in both transformer and VAE (and possibly a different file) because of some common parts shared, let's keep the implementation in transformer file, and import it in the vae. If there's no common attention processor, let's keep the implementation respectively in transformer or vae.
• If some specific layers are shared between transformer and vae (for example, GLUMB convolution in Sana), let's keep the implementation in transformer file too and import where required.
• Let's create dedicated RoPE classes for each implementation. Any concerns about speed due to recreating the embeddings every inference step can be addressed by caching. Something as simple as functools.cache works here if we make the rope calculation forward dependant on just the num_frames, height, width. But a save hook would work as well. This - we can look into a bit later

WDYT?

[yiyixuxu]

I think it will automtically upcast, no? here in hidden_states = hidden_states * scale if the hidden_states is in float32 an scale is in bfloat16, the operation will be in float32 and we do not need to explicitly upcast scale, no?

this is non-merge blocking! just curious if I misunderstood something

    def forward(self, hidden_states, scale=None):
        ....
        if scale is not None:
            hidden_states = hidden_states * scale

        hidden_states = hidden_states.to(hidden_states_dtype)

        return hidden_states

[DN6]

Failing LoRA tests are unrelated. Seem to be timing out on SDXL tests. Merging.