Add Photon model and pipeline support #12456
Conversation
This commit adds support for the Photon image generation model: - PhotonTransformer2DModel: Core transformer architecture - PhotonPipeline: Text-to-image generation pipeline - Attention processor updates for Photon-specific attention mechanism - Conversion script for loading Photon checkpoints - Documentation and tests
| print("✓ Created scheduler config") | ||
|
|
||
|
|
||
| def download_and_save_vae(vae_type: str, output_path: str): |
There was a problem hiding this comment.
I'm not sure on this one: I'm saving the VAE weights while they are already available on the Hub (Flux VAE and DC-AE).
Is there a way to avoid storing them and instead look directly for the original ones?
There was a problem hiding this comment.
For now, it's okay to keep this as is. This way, everything is under the same model repo.
| print(f"✓ Saved VAE to {vae_path}") | ||
|
|
||
|
|
||
| def download_and_save_text_encoder(output_path: str): |
There was a problem hiding this comment.
Same here for the Text Encoder.
| print("✓ Created scheduler config") | ||
|
|
||
|
|
||
| def download_and_save_vae(vae_type: str, output_path: str): |
There was a problem hiding this comment.
For now, it's okay to keep this as is. This way, everything is under the same model repo.
| from einops import rearrange | ||
| from einops.layers.torch import Rearrange |
There was a problem hiding this comment.
We need to get rid of the einops dependency and use native PyTorch ops here.
There was a problem hiding this comment.
I changed it for native Pytorch. Out of curiosity why do you recommend avoiding using einops?
There was a problem hiding this comment.
We try to avoid additional dependencies especially when things can be done in native PyTorch.
| return xq_out.reshape(*xq.shape).type_as(xq) | ||
|
|
||
|
|
||
| class EmbedND(nn.Module): |
There was a problem hiding this comment.
Does this share similarity with Flux?
There was a problem hiding this comment.
Yes it comes from the BFL original implementation.
I tried to modify and use the logic from transformer_flux.py but I didn't manage to make it work without heavy changes and additional complexity.
I added a comment to explicitely say that it come from there. Is it OK for you or do you want me to continue trying to use the code from transformer_flux.py?
There was a problem hiding this comment.
Oh okay. Then it's fine to keep it here. I would maybe rename it to PhotoEmbedND and leave a note that it's inspired from Flux. WDYT?
| - `sample` (`torch.Tensor`): Output latent image of shape `(B, C, H, W)`. | ||
| """ | ||
| if attention_kwargs is not None: |
There was a problem hiding this comment.
Could we unify the structure of this block similar to how it's done in QwenImage, for example (barring the bits related to ControlNet, of course)?
There was a problem hiding this comment.
I moved all the logic in this block.
3 participants
This commit adds support for the Photon image generation model:
What does this PR do?
Fixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
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.