Merge branch 'main' of https://github.com/coryMosaicML/diffusion into…

… aspect-ratio-buckets-betas-writer

diffusion/evaluation/clean_fid_eval.py

@@ -9,11 +9,10 @@
 
 import clip
 import torch
-import wandb
 from cleanfid import fid
 from composer import ComposerModel, Trainer
 from composer.core import get_precision_context
-from composer.loggers import LoggerDestination, WandBLogger
+from composer.loggers import LoggerDestination
 from composer.utils import dist
 from torch.utils.data import DataLoader
 from torchmetrics.multimodal import CLIPScore
@@ -91,19 +90,15 @@ def __init__(self,
         self.additional_generate_kwargs = additional_generate_kwargs if additional_generate_kwargs is not None else {}
         self.sdxl = model.sdxl
 
-        # Init loggers
-        if self.loggers and dist.get_local_rank() == 0:
-            for logger in self.loggers:
-                if isinstance(logger, WandBLogger):
-                    wandb.init(**logger._init_kwargs)
-
         # Load the model
-        Trainer(model=self.model,
-                load_path=self.load_path,
-                load_weights_only=True,
-                load_strict_model_weights=load_strict_model_weights,
-                eval_dataloader=self.eval_dataloader,
-                seed=self.seed)
+        trainer = Trainer(model=self.model,
+                          load_path=self.load_path,
+                          load_weights_only=True,
+                          load_strict_model_weights=load_strict_model_weights,
+                          eval_dataloader=self.eval_dataloader,
+                          seed=self.seed,
+                          loggers=self.loggers)
+        self.trainer = trainer
 
         # Move CLIP metric to device
         self.device = dist.get_local_rank()

diffusion/models/models.py

@@ -10,6 +10,7 @@
 import torch
 from composer.devices import DeviceGPU
 from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, EulerDiscreteScheduler, UNet2DConditionModel
+from peft import LoraConfig
 from torchmetrics import MeanSquaredError
 from transformers import CLIPTextModel, CLIPTokenizer, PretrainedConfig
 
@@ -67,6 +68,8 @@ def stable_diffusion_2(
     fsdp: bool = True,
     clip_qkv: Optional[float] = None,
     use_xformers: bool = True,
+    lora_rank: Optional[int] = None,
+    lora_alpha: Optional[int] = None,
 ):
     """Stable diffusion v2 training setup.
 
@@ -108,6 +111,8 @@ def stable_diffusion_2(
         fsdp (bool): Whether to use FSDP. Defaults to True.
         clip_qkv (float, optional): If not None, clip the qkv values to this value. Defaults to None.
         use_xformers (bool): Whether to use xformers for attention. Defaults to True.
+        lora_rank (int, optional): If not None, the rank to use for LoRA finetuning. Defaults to None.
+        lora_alpha (int, optional): If not None, the alpha to use for LoRA finetuning. Defaults to None.
     """
     latent_mean, latent_std = _parse_latent_statistics(latent_mean), _parse_latent_statistics(latent_std)
 
@@ -215,6 +220,40 @@ def stable_diffusion_2(
         mask_pad_tokens=mask_pad_tokens,
         fsdp=fsdp,
     )
+    if lora_rank is not None:
+        assert lora_alpha is not None
+        model.unet.requires_grad_(False)
+        for param in model.unet.parameters():
+            param.requires_grad_(False)
+
+        unet_lora_config = LoraConfig(
+            r=lora_rank,
+            lora_alpha=lora_alpha,
+            init_lora_weights='gaussian',
+            target_modules=['to_k', 'to_q', 'to_v', 'to_out.0'],
+        )
+        model.unet.add_adapter(unet_lora_config)
+        model.unet._fsdp_wrap = True
+        if hasattr(model.unet, 'mid_block') and model.unet.mid_block is not None:
+            for attention in model.unet.mid_block.attentions:
+                attention._fsdp_wrap = True
+            for resnet in model.unet.mid_block.resnets:
+                resnet._fsdp_wrap = True
+        for block in model.unet.up_blocks:
+            if hasattr(block, 'attentions'):
+                for attention in block.attentions:
+                    attention._fsdp_wrap = True
+            if hasattr(block, 'resnets'):
+                for resnet in block.resnets:
+                    resnet._fsdp_wrap = True
+        for block in model.unet.down_blocks:
+            if hasattr(block, 'attentions'):
+                for attention in block.attentions:
+                    attention._fsdp_wrap = True
+            if hasattr(block, 'resnets'):
+                for resnet in block.resnets:
+                    resnet._fsdp_wrap = True
+
     if torch.cuda.is_available():
         model = DeviceGPU().module_to_device(model)
         if is_xformers_installed and use_xformers:
@@ -264,6 +303,8 @@ def stable_diffusion_xl(
     fsdp: bool = True,
     clip_qkv: Optional[float] = None,
     use_xformers: bool = True,
+    lora_rank: Optional[int] = None,
+    lora_alpha: Optional[int] = None,
 ):
     """Stable diffusion 2 training setup + SDXL UNet and VAE.
 
@@ -319,6 +360,8 @@ def stable_diffusion_xl(
         clip_qkv (float, optional): If not None, clip the qkv values to this value. Improves stability of training.
             Default: ``None``.
         use_xformers (bool): Whether to use xformers for attention. Defaults to True.
+        lora_rank (int, optional): If not None, the rank to use for LoRA finetuning. Defaults to None.
+        lora_alpha (int, optional): If not None, the alpha to use for LoRA finetuning. Defaults to None.
     """
     latent_mean, latent_std = _parse_latent_statistics(latent_mean), _parse_latent_statistics(latent_std)
 
@@ -485,6 +528,40 @@ def stable_diffusion_xl(
         fsdp=fsdp,
         sdxl=True,
     )
+
+    if lora_rank is not None:
+        assert lora_alpha is not None
+        model.unet.requires_grad_(False)
+        for param in model.unet.parameters():
+            param.requires_grad_(False)
+
+        unet_lora_config = LoraConfig(
+            r=lora_rank,
+            lora_alpha=lora_alpha,
+            init_lora_weights='gaussian',
+            target_modules=['to_k', 'to_q', 'to_v', 'to_out.0'],
+        )
+        model.unet.add_adapter(unet_lora_config)
+        model.unet._fsdp_wrap = True
+        if hasattr(model.unet, 'mid_block') and model.unet.mid_block is not None:
+            for attention in model.unet.mid_block.attentions:
+                attention._fsdp_wrap = True
+            for resnet in model.unet.mid_block.resnets:
+                resnet._fsdp_wrap = True
+        for block in model.unet.up_blocks:
+            if hasattr(block, 'attentions'):
+                for attention in block.attentions:
+                    attention._fsdp_wrap = True
+            if hasattr(block, 'resnets'):
+                for resnet in block.resnets:
+                    resnet._fsdp_wrap = True
+        for block in model.unet.down_blocks:
+            if hasattr(block, 'attentions'):
+                for attention in block.attentions:
+                    attention._fsdp_wrap = True
+            if hasattr(block, 'resnets'):
+                for resnet in block.resnets:
+                    resnet._fsdp_wrap = True
     if torch.cuda.is_available():
         model = DeviceGPU().module_to_device(model)
         if is_xformers_installed and use_xformers:

setup.py

@@ -10,7 +10,7 @@
     'diffusers[torch]==0.26.3', 'transformers[torch]==4.38.2', 'huggingface_hub==0.21.2', 'wandb==0.16.3',
     'xformers==0.0.23.post1', 'triton==2.1.0', 'torchmetrics[image]==1.3.1', 'lpips==0.1.4', 'clean-fid==0.1.35',
     'clip@git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33', 'gradio==4.19.2',
-    'datasets==2.19.2'
+    'datasets==2.19.2', 'peft==0.12.0'
 ]
 
 extras_require = {}