Add Flux and Flux Controlnet Support to Diffusion folder by Victor49152 · Pull Request #11794 · NVIDIA-NeMo/NeMo

Victor49152 · 2025-01-08T23:32:06Z

What does this PR do ?

Add Flux and Flux Controlnet Support

Collection: [DIFFUSION]

torchrun flux_controlnet_infer.py --num_joint_layers 4 --num_single_layers 0 --flux_ckpt /ckpts/nemo_flux_transformer.safetensors --controlnet_ckpt /ckpts/nemo_flux_controlnet_transformer.safetensors --control_image /ckpts/FLUX.1-controlnet-lineart-promeai/images/example-control.jpg

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
[] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

…oising loop. Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

…anm/diffusion # Conflicts: # nemo/collections/diffusion/flux/pipeline.py

# Conflicts: # nemo/collections/diffusion/__init__.py # nemo/collections/diffusion/vae/__init__.py

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

github-actions · 2025-01-21T17:40:46Z

beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base.

Your code was analyzed with PyLint. The following annotations have been identified:

************* Module nemo.collections.diffusion.models.flux_controlnet.model
nemo/collections/diffusion/models/flux_controlnet/model.py:352:0: C0301: Line too long (120/119) (line-too-long)
************* Module nemo.collections.diffusion.encoders.conditioner
nemo/collections/diffusion/encoders/conditioner.py:23:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/encoders/conditioner.py:170:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.models.dit.dit_attention
nemo/collections/diffusion/models/dit/dit_attention.py:380:0: C0301: Line too long (122/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_attention.py:382:0: C0301: Line too long (122/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_attention.py:15:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.models.dit.dit_layer_spec
nemo/collections/diffusion/models/dit/dit_layer_spec.py:54:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:60:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:66:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:632:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:21:0: W0611: Unused rearrange imported from einops (unused-import)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:22:0: W0611: Unused jit_fuser imported from megatron.core.jit (unused-import)
************* Module nemo.collections.diffusion.models.flux.layers
nemo/collections/diffusion/models/flux/layers.py:131:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.models.flux.model
nemo/collections/diffusion/models/flux/model.py:720:0: C0301: Line too long (122/119) (line-too-long)
nemo/collections/diffusion/models/flux/model.py:74:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:107:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:113:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/flux/model.py:120:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.models.flux.pipeline
nemo/collections/diffusion/models/flux/pipeline.py:385:0: C0301: Line too long (120/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:492:0: C0301: Line too long (120/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:494:0: C0301: Line too long (120/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:800:0: C0301: Line too long (134/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:839:0: C0301: Line too long (125/119) (line-too-long)
nemo/collections/diffusion/models/flux/pipeline.py:650:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.utils.flux_pipeline_utils
nemo/collections/diffusion/utils/flux_pipeline_utils.py:15:0: W0611: Unused dataclass imported from dataclasses (unused-import)
nemo/collections/diffusion/utils/flux_pipeline_utils.py:17:0: W0611: Unused import torch (unused-import)
************* Module nemo.collections.diffusion.utils.mcore_parallel_utils
nemo/collections/diffusion/utils/mcore_parallel_utils.py:26:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.diffusion.vae.autoencoder
nemo/collections/diffusion/vae/autoencoder.py:26:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/vae/autoencoder.py:48:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/vae/autoencoder.py:149:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/vae/autoencoder.py:264:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/vae/autoencoder.py:279:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.lightning._strategy_lib
nemo/lightning/_strategy_lib.py:586:0: C0301: Line too long (130/119) (line-too-long)
nemo/lightning/_strategy_lib.py:35:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/_strategy_lib.py:36:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:141:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:168:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:204:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:517:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:611:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:624:4: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.lightning.megatron_parallel
nemo/lightning/megatron_parallel.py:245:0: C0301: Line too long (127/119) (line-too-long)
nemo/lightning/megatron_parallel.py:246:0: C0301: Line too long (140/119) (line-too-long)
nemo/lightning/megatron_parallel.py:247:0: C0301: Line too long (130/119) (line-too-long)
nemo/lightning/megatron_parallel.py:554:0: C0301: Line too long (129/119) (line-too-long)
nemo/lightning/megatron_parallel.py:561:0: C0301: Line too long (135/119) (line-too-long)
nemo/lightning/megatron_parallel.py:849:0: C0301: Line too long (137/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1079:0: C0301: Line too long (136/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1652:0: C0301: Line too long (128/119) (line-too-long)
nemo/lightning/megatron_parallel.py:1691:0: C0301: Line too long (146/119) (line-too-long)
nemo/lightning/megatron_parallel.py:71:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:72:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:74:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:109:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:113:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:313:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:337:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:363:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:389:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:525:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:569:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:573:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:639:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:674:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:680:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:686:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:693:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:700:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:734:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:742:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:758:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:785:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:797:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:819:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1345:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1520:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1526:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1532:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1536:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1541:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1546:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1574:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1620:8: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1642:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1715:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:1761:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1775:0: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.lightning.pytorch.optim.base
nemo/lightning/pytorch/optim/base.py:166:4: C0116: Missing function or method docstring (missing-function-docstring)

-----------------------------------
Your code has been rated at 9.72/10

Mitigation guide:

Add sensible and useful docstrings to functions and methods
For trivial methods like getter/setters, consider adding # pylint: disable=C0116 inside the function itself
To disable multiple functions/methods at once, put a # pylint: disable=C0116 before the first and a # pylint: enable=C0116 after the last.

By applying these rules, we reduce the occurance of this message in future.

Thank you for improving NeMo's documentation!

github-actions · 2025-01-21T21:32:26Z

[🤖]: Hi @Victor49152 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g

* Vae added and matched flux checkpoint Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux model added. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Copying FlowMatchEulerScheduler over Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * WIP: Start to test the pipeline forward pass Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Vae added and matched flux checkpoint Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Inference pipeline runs with offloading function Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Start to test image generation Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Decoding with VAE part has been verified. Still need to check the denoising loop. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * The inference pipeline is verified. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add arg parsers and refactoring Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Tested on multi batch sizes and prompts. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add headers Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Renaming Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Move shceduler to sampler folder Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Merging folders. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Tested after path changing. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Move MMDIT block to NeMo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Add joint attention and single attention to NeMo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Joint attention updated Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Remove redundant importing Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Refactor to inherit megatron module Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Adding mockdata Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * DDP training works Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added flux controlnet training components while not tested yet Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux training with DDP tested on 1 GPU Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux and controlnet now could train on precached mode. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Custom FSDP path added to megatron parallel. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Bug fix Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * A hacky way to wrap frozen flux into FSDP to reproduce illegal memory issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Typo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Bypass the no grad issue when no single layers exists Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * A hacky way to wrap frozen flux into FSDP to reproduce illegal memory issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Let the flux model's dtype autocast before FSDP wrapping * fix RuntimeError: "Output 0 of SliceBackward0 is a view and is being modified inplace..." * Add a wrapper to flux controlnet so they are all wrapped into FSDP automatically Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Get rid of concat op in flux single transformer Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Get rid of concat op in flux single transformer Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * single block attention.linear_proj.bias must not require grads after refactoring Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * use cpu initialization to avoid OOM Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Set up flux training script with tp Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * SDXL fid image generation script updated. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Mcore self attention API changed Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add a dummy task encoder for raw image inputs Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Support loading crudedataset via energon dataloader Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Default save last to True Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add controlnet inference pipeline Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add controlnet inference script Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image resize mode update Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Remove unnecessary bias to avoid sharding issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Handle MCore custom fsdp checkpoint load (#11621) * general handle custom_fsdp checkpoint load * Apply isort and black reformatting Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> * Apply isort and black reformatting Signed-off-by: artbataev <artbataev@users.noreply.github.com> --------- Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> Signed-off-by: artbataev <artbataev@users.noreply.github.com> Co-authored-by: shjwudp <shjwudp@users.noreply.github.com> Co-authored-by: artbataev <artbataev@users.noreply.github.com> * Checkpoint naming Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image logger WIP Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image logger works fine Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * save hint and output to image logger. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Update flux controlnet training step Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add model connector and try to load from dist ckpt but failed. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Renaming and refactoring submodel configs for nemo run compatibility Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Nemo run script works for basic testing recipe Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added tp2 training factory Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added convergence recipe Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added flux training scripts Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Inference script tested Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Controlnet inference script tested Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Moving scripts to correct folder and modify headers Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Doc strings update Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * pylint correction Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Add import guard since custom fsdp is not merged to mcore yet Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add copy right headers and correct code check Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Code Scan Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Minor fix Signed-off-by: mingyuanm <mingyuanm@nvidia.com> --------- Signed-off-by: mingyuanm <mingyuanm@nvidia.com> Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> Signed-off-by: artbataev <artbataev@users.noreply.github.com> Co-authored-by: Victor49152 <Victor49152@users.noreply.github.com> Co-authored-by: jianbinc <shjwudp@gmail.com> Co-authored-by: shjwudp <shjwudp@users.noreply.github.com> Co-authored-by: artbataev <artbataev@users.noreply.github.com> Signed-off-by: Abhinav Garg <abhgarg@nvidia.com>

…#11794) * Vae added and matched flux checkpoint Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux model added. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Copying FlowMatchEulerScheduler over Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * WIP: Start to test the pipeline forward pass Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Vae added and matched flux checkpoint Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Inference pipeline runs with offloading function Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Start to test image generation Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Decoding with VAE part has been verified. Still need to check the denoising loop. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * The inference pipeline is verified. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add arg parsers and refactoring Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Tested on multi batch sizes and prompts. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add headers Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Renaming Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Move shceduler to sampler folder Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Merging folders. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Tested after path changing. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Move MMDIT block to NeMo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Add joint attention and single attention to NeMo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Joint attention updated Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Remove redundant importing Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Refactor to inherit megatron module Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Adding mockdata Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * DDP training works Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added flux controlnet training components while not tested yet Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux training with DDP tested on 1 GPU Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Flux and controlnet now could train on precached mode. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Custom FSDP path added to megatron parallel. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Bug fix Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * A hacky way to wrap frozen flux into FSDP to reproduce illegal memory issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Typo Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Bypass the no grad issue when no single layers exists Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * A hacky way to wrap frozen flux into FSDP to reproduce illegal memory issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Let the flux model's dtype autocast before FSDP wrapping * fix RuntimeError: "Output 0 of SliceBackward0 is a view and is being modified inplace..." * Add a wrapper to flux controlnet so they are all wrapped into FSDP automatically Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Get rid of concat op in flux single transformer Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Get rid of concat op in flux single transformer Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * single block attention.linear_proj.bias must not require grads after refactoring Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * use cpu initialization to avoid OOM Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Set up flux training script with tp Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * SDXL fid image generation script updated. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Mcore self attention API changed Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add a dummy task encoder for raw image inputs Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Support loading crudedataset via energon dataloader Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Default save last to True Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add controlnet inference pipeline Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add controlnet inference script Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image resize mode update Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Remove unnecessary bias to avoid sharding issue. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Handle MCore custom fsdp checkpoint load (NVIDIA-NeMo#11621) * general handle custom_fsdp checkpoint load * Apply isort and black reformatting Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> * Apply isort and black reformatting Signed-off-by: artbataev <artbataev@users.noreply.github.com> --------- Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> Signed-off-by: artbataev <artbataev@users.noreply.github.com> Co-authored-by: shjwudp <shjwudp@users.noreply.github.com> Co-authored-by: artbataev <artbataev@users.noreply.github.com> * Checkpoint naming Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image logger WIP Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Image logger works fine Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * save hint and output to image logger. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Update flux controlnet training step Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add model connector and try to load from dist ckpt but failed. Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Renaming and refactoring submodel configs for nemo run compatibility Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Nemo run script works for basic testing recipe Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added tp2 training factory Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added convergence recipe Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Added flux training scripts Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Inference script tested Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Controlnet inference script tested Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Moving scripts to correct folder and modify headers Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Doc strings update Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * pylint correction Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Add import guard since custom fsdp is not merged to mcore yet Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Add copy right headers and correct code check Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Apply isort and black reformatting Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> * Code Scan Signed-off-by: mingyuanm <mingyuanm@nvidia.com> * Minor fix Signed-off-by: mingyuanm <mingyuanm@nvidia.com> --------- Signed-off-by: mingyuanm <mingyuanm@nvidia.com> Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com> Signed-off-by: shjwudp <shjwudp@users.noreply.github.com> Signed-off-by: artbataev <artbataev@users.noreply.github.com> Co-authored-by: Victor49152 <Victor49152@users.noreply.github.com> Co-authored-by: jianbinc <shjwudp@gmail.com> Co-authored-by: shjwudp <shjwudp@users.noreply.github.com> Co-authored-by: artbataev <artbataev@users.noreply.github.com> Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>

Victor49152 and others added 30 commits September 4, 2024 15:25

Vae added and matched flux checkpoint

0e3c818

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Flux model added.

8c9c56f

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Copying FlowMatchEulerScheduler over

9a304dc

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

WIP: Start to test the pipeline forward pass

73c714d

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Vae added and matched flux checkpoint

f4d7747

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Inference pipeline runs with offloading function

6e4de91

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Start to test image generation

2cb67f2

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Decoding with VAE part has been verified. Still need to check the den…

c18cf60

…oising loop. Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

The inference pipeline is verified.

072ce16

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Add arg parsers and refactoring

b4d281f

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Tested on multi batch sizes and prompts.

7d27534

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Add headers

6d2da09

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

db43ec7

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Renaming

597a646

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Merge remote-tracking branch 'origin/mingyuanm/diffusion' into mingyu…

d2bfbc3

…anm/diffusion # Conflicts: # nemo/collections/diffusion/flux/pipeline.py

Merge branch 'refs/heads/main' into mingyuanm/diffusion

7894f2c

# Conflicts: # nemo/collections/diffusion/__init__.py # nemo/collections/diffusion/vae/__init__.py

Move shceduler to sampler folder

6fb7433

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Merging folders.

f4cf498

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

756b8ee

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Tested after path changing.

73e2099

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

aec7a13

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Move MMDIT block to NeMo

15db8ad

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

6801903

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Add joint attention and single attention to NeMo

7d34b30

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

2bf20e1

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Joint attention updated

d78c682

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Apply isort and black reformatting

fbd6987

Signed-off-by: Victor49152 <Victor49152@users.noreply.github.com>

Remove redundant importing

aa9df2a

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Refactor to inherit megatron module

ae18bb6

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Adding mockdata

94b1a3d

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Victor49152 dismissed yaoyu-33’s stale review via f25405c January 15, 2025 20:04

Victor49152 enabled auto-merge (squash) January 15, 2025 20:06

yaoyu-33 previously approved these changes Jan 15, 2025

View reviewed changes

Victor49152 added Run CICD CI and removed Run CICD labels Jan 16, 2025

Merge branch 'refs/heads/main' into mingyuanm/flux_controlnet

896ac09

Victor49152 dismissed yaoyu-33’s stale review via 896ac09 January 16, 2025 18:08

github-actions bot removed the CI label Jan 16, 2025

Victor49152 added the Run CICD label Jan 16, 2025

Merge branch 'main' into mingyuanm/flux_controlnet

0a1592b

Victor49152 added Run CICD and removed Run CICD labels Jan 17, 2025

Merge branch 'main' into mingyuanm/flux_controlnet

3952667

Victor49152 added Run CICD and removed Run CICD labels Jan 17, 2025

Merge branch 'main' into mingyuanm/flux_controlnet

45fbb1e

Victor49152 added Run CICD and removed Run CICD labels Jan 18, 2025

Update megatron fsdp guard for importing errors

e7a151c

Signed-off-by: mingyuanm <mingyuanm@nvidia.com>

Victor49152 added Run CICD and removed Run CICD labels Jan 21, 2025

yaoyu-33 approved these changes Jan 21, 2025

View reviewed changes

Victor49152 merged commit 066e4b4 into main Jan 21, 2025
211 of 213 checks passed

Victor49152 deleted the mingyuanm/flux_controlnet branch January 21, 2025 22:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Flux and Flux Controlnet Support to Diffusion folder#11794

Add Flux and Flux Controlnet Support to Diffusion folder#11794
Victor49152 merged 94 commits intomainfrom
mingyuanm/flux_controlnet

Victor49152 commented Jan 8, 2025

Uh oh!

github-actions bot commented Jan 21, 2025

Uh oh!

github-actions bot commented Jan 21, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

Victor49152 commented Jan 8, 2025

What does this PR do ?

Before your PR is "Ready for review"

Uh oh!

github-actions bot commented Jan 21, 2025

Uh oh!

github-actions bot commented Jan 21, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants