Add Flux and Flux Controlnet Support to Diffusion folder #11794

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 <[email protected]>

…oising loop. Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

…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 <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

# Conflicts: # nemo/collections/diffusion/vae/autoencoder.py # nemo/lightning/megatron_parallel.py # scripts/dit/dit_train.py

Signed-off-by: mingyuanm <[email protected]>

nemo/collections/diffusion/models/flux/model.py

nemo/collections/diffusion/models/flux/pipeline.py

nemo/collections/diffusion/models/flux_controlnet/model.py

+        Saves the inference results together with the hint image to log folder.
+        '''
+        logging.info("Start validation step")
+        from nemo.collections.diffusion.models.flux.pipeline import FluxControlNetInferencePipeline


nemo/lightning/_strategy_lib.py

+        from megatron.core.distributed.custom_fsdp import FullyShardedDataParallel
+
+        use_custom_fsdp = True
+    except:


nemo/lightning/_strategy_lib.py

+        use_custom_fsdp = True
+    except:
+        use_custom_fsdp = False
+        pass


nemo/lightning/megatron_parallel.py

+            # if not self._cpu:
+            #     model_module.cuda(torch.cuda.current_device())


scripts/flux/flux_controlnet_training.py

scripts/flux/flux_training.py

Signed-off-by: mingyuanm <[email protected]>

Signed-off-by: Victor49152 <[email protected]>

Victor49152 · 2025-01-09T19:21:02Z

@ethanhe42 I have moved the dit scripts into scrips folder as a NeMo 2 convention. Please take a look.

github-actions · 2025-01-10T17:12:12Z

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:587: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:612:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/_strategy_lib.py:625:4: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.lightning.megatron_parallel
nemo/lightning/megatron_parallel.py:243:0: C0301: Line too long (127/119) (line-too-long)
nemo/lightning/megatron_parallel.py:244:0: C0301: Line too long (140/119) (line-too-long)
nemo/lightning/megatron_parallel.py:245: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:69:0: C0115: Missing class docstring (missing-class-docstring)
nemo/lightning/megatron_parallel.py:70:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:72:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:107:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:111:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:311:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:335:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:361:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:387:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:523: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:1758:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/lightning/megatron_parallel.py:1772: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!

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

Vae added and matched flux checkpoint

0e3c818

Signed-off-by: mingyuanm <[email protected]>

Flux model added.

8c9c56f

Signed-off-by: mingyuanm <[email protected]>

Copying FlowMatchEulerScheduler over

9a304dc

Signed-off-by: mingyuanm <[email protected]>

WIP: Start to test the pipeline forward pass

73c714d

Signed-off-by: mingyuanm <[email protected]>

Vae added and matched flux checkpoint

f4d7747

Signed-off-by: mingyuanm <[email protected]>

Inference pipeline runs with offloading function

6e4de91

Signed-off-by: mingyuanm <[email protected]>

Start to test image generation

2cb67f2

Signed-off-by: mingyuanm <[email protected]>

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

c18cf60

…oising loop. Signed-off-by: mingyuanm <[email protected]>

The inference pipeline is verified.

072ce16

Signed-off-by: mingyuanm <[email protected]>

Add arg parsers and refactoring

b4d281f

Signed-off-by: mingyuanm <[email protected]>

Tested on multi batch sizes and prompts.

7d27534

Signed-off-by: mingyuanm <[email protected]>

Add headers

6d2da09

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

db43ec7

Signed-off-by: Victor49152 <[email protected]>

Renaming

597a646

Signed-off-by: mingyuanm <[email protected]>

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 <[email protected]>

Merging folders.

f4cf498

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

756b8ee

Signed-off-by: Victor49152 <[email protected]>

Tested after path changing.

73e2099

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

aec7a13

Signed-off-by: Victor49152 <[email protected]>

Move MMDIT block to NeMo

15db8ad

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

6801903

Signed-off-by: Victor49152 <[email protected]>

Add joint attention and single attention to NeMo

7d34b30

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

2bf20e1

Signed-off-by: Victor49152 <[email protected]>

Joint attention updated

d78c682

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

fbd6987

Signed-off-by: Victor49152 <[email protected]>

Remove redundant importing

aa9df2a

Signed-off-by: mingyuanm <[email protected]>

Refactor to inherit megatron module

ae18bb6

Signed-off-by: mingyuanm <[email protected]>

Adding mockdata

94b1a3d

Signed-off-by: mingyuanm <[email protected]>

Victor49152 added 7 commits January 5, 2025 20:18

Nemo run script works for basic testing recipe

2dfab9e

Signed-off-by: mingyuanm <[email protected]>

Added tp2 training factory

0d874f5

Signed-off-by: mingyuanm <[email protected]>

Added convergence recipe

ec3154e

Signed-off-by: mingyuanm <[email protected]>

Added flux training scripts

0aab7dc

Signed-off-by: mingyuanm <[email protected]>

Inference script tested

5053c41

Signed-off-by: mingyuanm <[email protected]>

Controlnet inference script tested

0d73d69

Signed-off-by: mingyuanm <[email protected]>

Moving scripts to correct folder and modify headers

b016614

Signed-off-by: mingyuanm <[email protected]>

Victor49152 requested review from ethanhe42 and yaoyu-33 January 8, 2025 23:32

github-actions bot added the Multi Modal label Jan 8, 2025

Victor49152 self-assigned this Jan 8, 2025

Victor49152 and others added 7 commits January 8, 2025 23:32

Apply isort and black reformatting

848bd27

Signed-off-by: Victor49152 <[email protected]>

Doc strings update

6129c82

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

790abdd

Signed-off-by: Victor49152 <[email protected]>

pylint correction

90a20e8

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

36bcf5c

Signed-off-by: Victor49152 <[email protected]>

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

36644c8

# Conflicts: # nemo/collections/diffusion/vae/autoencoder.py # nemo/lightning/megatron_parallel.py # scripts/dit/dit_train.py

Add import guard since custom fsdp is not merged to mcore yet

1ec2c16

Signed-off-by: mingyuanm <[email protected]>

Victor49152 added Run CICD PoR Major feature to be highlighted in release notes labels Jan 9, 2025

github-advanced-security bot found potential problems Jan 9, 2025

View reviewed changes

Victor49152 and others added 2 commits January 9, 2025 11:06

Add copy right headers and correct code check

6c2cbeb

Signed-off-by: mingyuanm <[email protected]>

Apply isort and black reformatting

20a52e3

Signed-off-by: Victor49152 <[email protected]>

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

Merge branch 'main' into mingyuanm/flux_controlnet

3dcdfe6

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 commented Jan 8, 2025

Victor49152 commented Jan 9, 2025

github-actions bot commented Jan 10, 2025

		# if not self._cpu:
		# model_module.cuda(torch.cuda.current_device())

Add Flux and Flux Controlnet Support to Diffusion folder #11794

Are you sure you want to change the base?

Add Flux and Flux Controlnet Support to Diffusion folder #11794

Conversation

Victor49152 commented Jan 8, 2025

What does this PR do ?

Before your PR is "Ready for review"

Victor49152 commented Jan 9, 2025

github-actions bot commented Jan 10, 2025