wip adding AMP and xformers to training code path

StableDiffuser.py

@@ -4,6 +4,7 @@ import torch
 from baukit import TraceDict
 from diffusers import StableDiffusionPipeline
 from PIL import Image
+from torch.cuda.amp import GradScaler
 from tqdm.auto import tqdm
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
@@ -34,18 +35,17 @@ class StableDiffuser(torch.nn.Module):
 
     def __init__(self,
                 scheduler='LMS',
-                 repo_id_or_path="CompVis/stable-diffusion-v1-4",
-                 variant='fp16'
-        ):
+                 repo_id_or_path="CompVis/stable-diffusion-v1-4"):
 
         super().__init__()
 
-        self.pipeline = StableDiffusionPipeline.from_pretrained(repo_id_or_path, variant=variant)
+        self.pipeline = StableDiffusionPipeline.from_pretrained(repo_id_or_path)
 
         self.vae = self.pipeline.vae
         self.unet = self.pipeline.unet
         self.tokenizer = self.pipeline.tokenizer
         self.text_encoder = self.pipeline.text_encoder
+        self.safety_checker = self.pipeline.safety_checker
 
         if scheduler == 'LMS':
             self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
@@ -57,8 +57,8 @@ class StableDiffuser(torch.nn.Module):
         self.eval()
 
     @property
-    def safety_checker(self):
-        return self.pipeline.safety_checker
+    def feature_extractor(self):
+        return self.pipeline.feature_extractor
 
     def get_noise(self, batch_size, width, height, generator=None):
         param = list(self.parameters())[0]
@@ -215,9 +215,9 @@ class StableDiffuser(torch.nn.Module):
                 self.safety_checker = self.safety_checker.float()
                 safety_checker_input = self.feature_extractor(images_steps[i], return_tensors="pt").to(latents_steps[0].device)
                 image, has_nsfw_concept = self.safety_checker(
-                    images=latents_steps[i].float().cpu().numpy(), clip_input=safety_checker_input.pixel_values.float()
+                    images=latents_steps[i], clip_input=safety_checker_input.pixel_values.float()
                 )
-                images_steps[i][0] = self.to_image(torch.from_numpy(image))[0]
+                images_steps[i][0] = self.to_image(image)[0]
 
         images_steps = list(zip(*images_steps))
 

app.py

@@ -1,8 +1,12 @@
 import gradio as gr
 import torch
 import os
+
+from diffusers.utils import is_xformers_available
+
 from finetuning import FineTunedModel
 from StableDiffuser import StableDiffuser
+from memory_efficiency import MemoryEfficiencyWrapper
 from train import train
     
 import os
@@ -158,8 +162,8 @@ class Demo:
                             info="Prompt corresponding to concept to erase"
                         )
 
-                        choices = ['ESD-x']
-                        if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40:
+                        choices = ['ESD-x', 'ESD-self']
+                        if torch.cuda.get_device_properties(0).total_memory * 1e-9 >= 40 or is_xformers_available():
                             choices.append('ESD-u')
                     
                         self.train_method_input = gr.Dropdown(
@@ -188,6 +192,12 @@ class Demo:
                             info='Learning rate used to train'
                         )
 
+                        with gr.Row():
+                            self.train_use_adamw8bit_input = gr.Checkbox(label="8bit AdamW", value=False)
+                            self.train_use_xformers_input = gr.Checkbox(label="xformers", value=True)
+                            self.train_use_amp_input = gr.Checkbox(label="AMP", value=True)
+                            #self.train_use_gradient_checkpointing_input = gr.Checkbox(label="Gradient checkpointing", value=True)
+
                     with gr.Column(scale=1):
 
                         self.train_status = gr.Button(value='', variant='primary', label='Status', interactive=False)
@@ -258,7 +268,11 @@ class Demo:
             self.train_method_input, 
             self.neg_guidance_input,
             self.iterations_input,
-            self.lr_input
+            self.lr_input,
+            self.train_use_adamw8bit_input,
+            self.train_use_xformers_input,
+            self.train_use_amp_input,
+            #self.train_use_gradient_checkpointing_input
         ],
         outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
         )
@@ -271,41 +285,43 @@ class Demo:
         outputs=[self.export_button]
         )
 
-    def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
+    def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr,
+              use_adamw8bit=True, use_xformers=True, use_amp=True, use_gradient_checkpointing=True,
+              pbar = gr.Progress(track_tqdm=True)):
 
         if self.training:
             return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
 
         if train_method == 'ESD-x':
-
             modules = ".*attn2$"
             frozen = []
 
         elif train_method == 'ESD-u':
-
             modules = "unet$"
             frozen = [".*attn2$", "unet.time_embedding$", "unet.conv_out$"]   
 
         elif train_method == 'ESD-self':
-
             modules = ".*attn1$"
             frozen = []
 
         randn = torch.randint(1, 10000000, (1,)).item()
 
-        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}.pt"
-
-        self.training = True
-
-        train(repo_id_or_path, img_size, prompt, modules, frozen, iterations, neg_guidance, lr, save_path)
-
-        self.training = False
+        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}_{train_method}_ng{neg_guidance}_lr{lr}_iter{iterations}.pt"
+        try:
+            self.training = True
+            train(repo_id_or_path, img_size, prompt, modules, frozen, iterations, neg_guidance, lr, save_path,
+                use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing)
+        finally:
+            self.training = False
 
         torch.cuda.empty_cache()
 
-        model_map['Custom'] = save_path
+        new_model_name = f'*new* {os.path.basename(save_path)}'
+        model_map[new_model_name] = save_path
 
-        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
+        return [gr.update(interactive=True, value='Train'), gr.update(value=f'Done Training! \n '
+                'Try your model ({new_model_name}) in the "Test" tab'), save_path,
+                gr.Dropdown.update(choices=list(model_map.keys()), value=new_model_name)]
 
     def export(self, model_name, base_repo_id_or_path, save_path, save_half):
         model_path = model_map[model_name]

finetuning.py

@@ -51,7 +51,6 @@ class FineTunedModel(torch.nn.Module):
 
     @classmethod
     def from_checkpoint(cls, model, checkpoint, frozen_modules=[]):
-
         if isinstance(checkpoint, str):
             checkpoint = torch.load(checkpoint)
 
@@ -64,33 +63,23 @@ class FineTunedModel(torch.nn.Module):
 
         
     def __enter__(self):
-
         for key, ft_module in self.ft_modules.items():
             util.set_module(self.model, key, ft_module)
 
     def __exit__(self, exc_type, exc_value, tb):
-
         for key, module in self.orig_modules.items():
             util.set_module(self.model, key, module)
 
     def parameters(self):
-
         parameters = []
-
         for ft_module in self.ft_modules.values():
-
             parameters.extend(list(ft_module.parameters()))
-
         return parameters
 
     def state_dict(self):
-
         state_dict = {key: module.state_dict() for key, module in self.ft_modules.items()}
-
         return state_dict
 
     def load_state_dict(self, state_dict):
-
         for key, sd in state_dict.items():
-            
             self.ft_modules[key].load_state_dict(sd)

memory_efficiency.py

@@ -0,0 +1,86 @@
+# adapted from EveryDream2Trainer
+import contextlib
+import traceback
+
+import torch
+from torch.cuda.amp import GradScaler
+
+from StableDiffuser import StableDiffuser
+
+
+class MemoryEfficiencyWrapper:
+
+    def __init__(self,
+                 diffuser: StableDiffuser,
+                 use_amp: bool,
+                 use_xformers: bool,
+                 use_gradient_checkpointing: bool):
+        self.diffuser = diffuser
+        self.is_sd1attn = diffuser.unet.config["attention_head_dim"] == [8, 8, 8, 8]
+        self.is_sd1attn = diffuser.unet.config["attention_head_dim"] == 8 or self.is_sd1attn
+
+        self.use_amp = use_amp
+        self.use_xformers = use_xformers
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+
+    def __enter__(self):
+        if self.use_gradient_checkpointing:
+            self.diffuser.unet.enable_gradient_checkpointing()
+            self.diffuser.text_encoder.gradient_checkpointing_enable()
+
+        if self.use_xformers:
+            if (self.use_amp and self.is_sd1attn) or (not self.is_sd1attn):
+                try:
+                    self.diffuser.unet.enable_xformers_memory_efficient_attention()
+                    print("Enabled xformers")
+                except Exception as ex:
+                    print("failed to load xformers, using attention slicing instead")
+                    self.diffuser.unet.set_attention_slice("auto")
+                    pass
+            elif (not self.amp and self.is_sd1attn):
+                print("AMP is disabled but model is SD1.X, using attention slicing instead of xformers")
+                self.diffuser.unet.set_attention_slice("auto")
+        else:
+            print("xformers disabled via arg, using attention slicing instead")
+            self.diffuser.unet.set_attention_slice("auto")
+
+        self.diffuser.vae = self.diffuser.vae.to(self.diffuser.vae.device, dtype=torch.float16 if self.use_amp else torch.float32)
+        self.diffuser.unet = self.diffuser.unet.to(self.diffuser.unet.device, dtype=torch.float32)
+
+        try:
+            # unet = torch.compile(unet)
+            # text_encoder = torch.compile(text_encoder)
+            # vae = torch.compile(vae)
+            torch.set_float32_matmul_precision('high')
+            torch.backends.cudnn.allow_tf32 = True
+            # logging.info("Successfully compiled models")
+        except Exception as ex:
+            print(f"Failed to compile model, continuing anyway, ex: {ex}")
+            pass
+
+        self.grad_scaler = GradScaler(
+            enabled=self.use_amp,
+            init_scale=2 ** 17.5,
+            growth_factor=2,
+            backoff_factor=1.0 / 2,
+            growth_interval=25,
+        )
+
+    def step(self, optimizer, loss):
+        self.grad_scaler.scale(loss).backward()
+        self.grad_scaler.step(optimizer)
+        self.grad_scaler.update()
+
+    def __exit__(self, exc_type, exc_value, tb):
+        if exc_type is not None:
+            traceback.print_exception(exc_type, exc_value, tb)
+        # return False # uncomment to pass exception through):
+        self.diffuser.unet.disable_gradient_checkpointing()
+        try:
+            self.diffuser.text_encoder.gradient_checkpointing_disable()
+        except AttributeError:
+            # self.diffuser.text_encoder is likely `del`eted
+            pass
+
+        self.diffuser.unet.disable_xformers_memory_efficient_attention()
+        self.diffuser.unet.set_attention_slice("auto")

requirements.txt

@@ -1,8 +1,11 @@
 gradio
-torch==1.13.1 --index-url https://download.pytorch.org/whl/cu118
-torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu118
+torch --index-url https://download.pytorch.org/whl/cu118
+torchvision --index-url https://download.pytorch.org/whl/cu118
 diffusers
 transformers
 accelerate
 scipy
 git+https://github.com/davidbau/baukit.git
+xformers
+bitsandbytes==0.38.1
+safetensors

train.py

@@ -3,68 +3,85 @@ from finetuning import FineTunedModel
 import torch
 from tqdm import tqdm
 
-def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path):
+from memory_efficiency import MemoryEfficiencyWrapper
+
+
+def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path,
+          use_adamw8bit=True, use_xformers=True, use_amp=True, use_gradient_checkpointing=False):
   
     nsteps = 50
 
     diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=repo_id_or_path).to('cuda')
-    diffuser.train()
 
-    finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
+    memory_efficiency_wrapper = MemoryEfficiencyWrapper(diffuser=diffuser, use_amp=use_amp, use_xformers=use_xformers,
+                                                        use_gradient_checkpointing=use_gradient_checkpointing )
+    with memory_efficiency_wrapper:
 
-    optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
-    criteria = torch.nn.MSELoss()
+        diffuser.train()
 
-    pbar = tqdm(range(iterations))
+        finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
 
-    with torch.no_grad():
+        if use_adamw8bit:
+            import bitsandbytes as bnb
+            optimizer = bnb.optim.AdamW8bit(finetuner.parameters(),
+                                            lr=lr,
+                                            betas=(0.9, 0.999),
+                                            weight_decay=0.010,
+                                            eps=1e-8
+                                            )
+        else:
+            optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
+        criteria = torch.nn.MSELoss()
 
-        neutral_text_embeddings = diffuser.get_text_embeddings([''],n_imgs=1)
-        positive_text_embeddings = diffuser.get_text_embeddings([prompt],n_imgs=1)
+        pbar = tqdm(range(iterations))
 
-    del diffuser.vae
-    del diffuser.text_encoder
-    del diffuser.tokenizer
+        with torch.no_grad():
 
-    torch.cuda.empty_cache()
+            neutral_text_embeddings = diffuser.get_text_embeddings([''],n_imgs=1)
+            positive_text_embeddings = diffuser.get_text_embeddings([prompt],n_imgs=1)
 
-    print(f"using img_size of {img_size}")
+        del diffuser.vae
+        del diffuser.text_encoder
+        del diffuser.tokenizer
 
-    for i in pbar:
-        with torch.no_grad():
-            diffuser.set_scheduler_timesteps(nsteps)
-            optimizer.zero_grad()
+        torch.cuda.empty_cache()
+
+        print(f"using img_size of {img_size}")
+
+        for i in pbar:
+            with torch.no_grad():
+                diffuser.set_scheduler_timesteps(nsteps)
+                optimizer.zero_grad()
 
-            iteration = torch.randint(1, nsteps - 1, (1,)).item()
-            latents = diffuser.get_initial_latents(1, width=img_size, height=img_size, n_prompts=1)
+                iteration = torch.randint(1, nsteps - 1, (1,)).item()
+                latents = diffuser.get_initial_latents(1, width=img_size, height=img_size, n_prompts=1)
+
+                with finetuner:
+                    latents_steps, _ = diffuser.diffusion(
+                        latents,
+                        positive_text_embeddings,
+                        start_iteration=0,
+                        end_iteration=iteration,
+                        guidance_scale=3,
+                        show_progress=False
+                    )
+
+                diffuser.set_scheduler_timesteps(1000)
+
+                iteration = int(iteration / nsteps * 1000)
+
+                positive_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+                neutral_latents = diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1)
 
             with finetuner:
-                latents_steps, _ = diffuser.diffusion(
-                    latents,
-                    positive_text_embeddings,
-                    start_iteration=0,
-                    end_iteration=iteration,
-                    guidance_scale=3, 
-                    show_progress=False
-                )
-
-            diffuser.set_scheduler_timesteps(1000)
-
-            iteration = int(iteration / nsteps * 1000)
-            
-            positive_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
-            neutral_latents = diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1)
-
-        with finetuner:
-            negative_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
-
-        positive_latents.requires_grad = False
-        neutral_latents.requires_grad = False
-
-        loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents))) #loss = criteria(e_n, e_0) works the best try 5000 epochs
-        
-        loss.backward()
-        optimizer.step()
+                negative_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+
+            positive_latents.requires_grad = False
+            neutral_latents.requires_grad = False
+
+            loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents))) #loss = criteria(e_n, e_0) works the best try 5000 epochs
+            memory_efficiency_wrapper.step(optimizer, loss)
+            optimizer.step()
 
     torch.save(finetuner.state_dict(), save_path)