Add Flux and Flux Controlnet Support to Diffusion folder

examples/multimodal/text_to_image/stable_diffusion/conf/sd_xl_fid_images.yaml

@@ -2,25 +2,15 @@ name: stable-diffusion-train
 
 fid:
   classifier_free_guidance:
-    - 1.5
-    - 2
-    - 3
-    - 4
-    - 5
-    - 6
     - 7
-    - 8
-  nnodes_per_cfg: 1
+  nnodes_per_cfg: 2
   ntasks_per_node: 8
   local_task_id: null
   num_images_to_eval: 30000
-  coco_captions_path: /coco2014/coco2014_val_sampled_30k/captions
-  coco_images_path: /coco2014/coco2014_val/images_256
+  coco_captions_path: /datasets/coco2014/coco2014_val_sampled_30k/captions
+  coco_images_path: /datasets/coco2014/coco2014_val/images_256
   save_path: output
 
-model:
-  restore_from_path:
-  is_legacy: False
 
 use_refiner: False
 use_fp16: False # use fp16 model weights
@@ -88,8 +78,128 @@ sampling:
     order: 4
 
 trainer:
-  devices: ${evaluation.fid.ntasks_per_node}
+  devices: ${fid.ntasks_per_node}
   num_nodes: 1
   accelerator: gpu
   precision: 32
   logger: False # logger provided by exp_manager
+
+
+model:
+  restore_from_path: null
+  is_legacy: False
+  scale_factor: 0.13025
+  disable_first_stage_autocast: True
+
+  fsdp: False
+  fsdp_set_buffer_dtype: null
+  fsdp_sharding_strategy: 'full'
+  use_cpu_initialization: True
+
+  optim:
+    name: fused_adam
+    lr: 1e-4
+    weight_decay: 0.0
+    betas:
+      - 0.9
+      - 0.999
+    sched:
+      name: WarmupHoldPolicy
+      warmup_steps: 10
+      hold_steps: 10000000000000 # Incredibly large value to hold the lr as constant
+
+  denoiser_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser.DiscreteDenoiser
+    num_idx: 1000
+
+    weighting_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser_weighting.EpsWeighting
+    scaling_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser_scaling.EpsScaling
+    discretization_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+  unet_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel.UNetModel
+    from_pretrained:
+    from_NeMo: True
+    adm_in_channels: 2816
+    num_classes: sequential
+    use_checkpoint: False
+    in_channels: 4
+    out_channels: 4
+    model_channels: 320
+    attention_resolutions: [ 4, 2 ]
+    num_res_blocks: 2
+    channel_mult: [ 1, 2, 4 ]
+    num_head_channels: 64
+    use_spatial_transformer: True
+    use_linear_in_transformer: True
+    transformer_depth: [ 1, 2, 10 ]  # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+    context_dim: 2048
+    image_size: 64 # unused
+    #    spatial_transformer_attn_type: softmax  #note: only default softmax is supported now
+    legacy: False
+    use_flash_attention: False
+
+  first_stage_config:
+    # _target_: nemo.collections.multimodal.models.stable_diffusion.ldm.autoencoder.AutoencoderKLInferenceWrapper
+    _target_: nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder.AutoencoderKLInferenceWrapper
+    from_pretrained:
+    from_NeMo: False
+    embed_dim: 4
+    monitor: val/rec_loss
+    ddconfig:
+      attn_type: vanilla
+      double_z: true
+      z_channels: 4
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [ 1, 2, 4, 4 ]
+      num_res_blocks: 2
+      attn_resolutions: [ ]
+      dropout: 0.0
+    lossconfig:
+      target: torch.nn.Identity
+
+  conditioner_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.GeneralConditioner
+    emb_models:
+      # crossattn cond
+      - is_trainable: False
+        input_key: txt
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenCLIPEmbedder
+          layer: hidden
+          layer_idx: 11
+      # crossattn and vector cond
+      - is_trainable: False
+        input_key: txt
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenOpenCLIPEmbedder2
+          arch: ViT-bigG-14
+          version: laion2b_s39b_b160k
+          freeze: True
+          layer: penultimate
+          always_return_pooled: True
+          legacy: False
+      # vector cond
+      - is_trainable: False
+        input_key: original_size_as_tuple
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two
+      # vector cond
+      - is_trainable: False
+        input_key: crop_coords_top_left
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two
+      # vector cond
+      - is_trainable: False
+        input_key: target_size_as_tuple
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two

examples/multimodal/text_to_image/stable_diffusion/conf/sd_xl_infer_v2.yaml

@@ -70,7 +70,7 @@ model:
   scale_factor: 0.13025
   disable_first_stage_autocast: True
   is_legacy: False
-  restore_from_path: ""
+  restore_from_path: null
 
   fsdp: False
   fsdp_set_buffer_dtype: null

examples/multimodal/text_to_image/stable_diffusion/generate_xl_fid_images.py

@@ -26,8 +26,9 @@
 from nemo.core.config import hydra_runner
 
 
-@hydra_runner(config_path='conf/stable_diffusion/conf', config_name='sd_xl_fid_images')
+@hydra_runner(config_path='conf', config_name='sd_xl_fid_images')
 def main(cfg):
+    # pylint: disable=C0116
     # Read configuration parameters
     nnodes_per_cfg = cfg.fid.nnodes_per_cfg
     ntasks_per_node = cfg.fid.ntasks_per_node