actually use AMP=3x speedup

StableDiffuser.py

@@ -4,7 +4,6 @@ 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
@@ -35,6 +34,7 @@ class StableDiffuser(torch.nn.Module):
 
     def __init__(self,
                 scheduler='LMS',
+                 keep_pipeline=False,
                  repo_id_or_path="CompVis/stable-diffusion-v1-4"):
 
         super().__init__()
@@ -46,6 +46,7 @@ class StableDiffuser(torch.nn.Module):
         self.tokenizer = self.pipeline.tokenizer
         self.text_encoder = self.pipeline.text_encoder
         self.safety_checker = self.pipeline.safety_checker
+        self.feature_extractor = self.pipeline.feature_extractor
 
         if scheduler == 'LMS':
             self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
@@ -55,10 +56,8 @@ class StableDiffuser(torch.nn.Module):
             self.scheduler = DDPMScheduler.from_pretrained(repo_id_or_path, subfolder="scheduler")
 
         self.eval()
-
-    @property
-    def feature_extractor(self):
-        return self.pipeline.feature_extractor
+        if not keep_pipeline:
+            del self.pipeline
 
     def get_noise(self, batch_size, width, height, generator=None):
         param = list(self.parameters())[0]
@@ -226,9 +225,6 @@ class StableDiffuser(torch.nn.Module):
 
         return images_steps
 
-    def save_pretrained(self, path, **kwargs):
-        self.pipeline.save_pretrained(path, **kwargs)
-
 
 if __name__ == '__main__':
 

app.py

@@ -162,9 +162,9 @@ class Demo:
                             info="Prompt corresponding to concept to erase"
                         )
 
-                        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')
+                        choices = ['ESD-x', 'ESD-self', 'ESD-u']
+                        #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(
                             choices=choices,
@@ -274,7 +274,7 @@ class Demo:
             self.train_use_amp_input,
             #self.train_use_gradient_checkpointing_input
         ],
-        outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
+        outputs=[self.train_button, self.train_status, self.download, self.model_dropdown]
         )
         self.export_button.click(self.export, inputs = [
             self.model_dropdown_export,
@@ -286,12 +286,19 @@ class Demo:
         )
 
     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,
+              use_adamw8bit=True, use_xformers=False, use_amp=False, use_gradient_checkpointing=False,
               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()]
 
+        print(f"Training {repo_id_or_path} at {img_size} to remove '{prompt}'.")
+        print(f"  {train_method}, negative guidance {neg_guidance}, lr {lr}, {iterations} iterations.")
+        print(f" {'✅' if use_gradient_checkpointing else '❌'} gradient checkpointing")
+        print(f" {'✅' if use_amp else '❌'} AMP")
+        print(f" {'✅' if use_xformers else '❌'} xformers")
+        print(f" {'✅' if use_adamw8bit else '❌'} 8-bit AdamW")
+
         if train_method == 'ESD-x':
             modules = ".*attn2$"
             frozen = []
@@ -319,20 +326,24 @@ class Demo:
         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=f'Done Training! \n '
-                'Try your model ({new_model_name}) in the "Test" tab'), save_path,
+        return [gr.update(interactive=True, value='Train'),
+                gr.update(value=f'Done Training! 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]
         checkpoint = torch.load(model_path)
-        self.diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=base_repo_id_or_path).to('cuda').eval()
-        finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval()
+        diffuser = StableDiffuser(scheduler='DDIM',
+                                       keep_pipeline=True,
+                                       repo_id_or_path=base_repo_id_or_path
+                                       ).eval()
+        finetuner = FineTunedModel.from_checkpoint(diffuser, checkpoint).eval()
         with finetuner:
             if save_half:
-                self.diffuser = self.diffuser.half()
-                self.diffuser.pipeline.to(torch.float16, torch_device=self.diffuser.device)
-            self.diffuser.save_pretrained(save_path)
+                diffuser = diffuser.half()
+                diffuser.pipeline.to(torch.float16, torch_device=diffuser.device)
+            diffuser.pipeline.save_pretrained(save_path)
 
 
     def inference(self, prompt, negative_prompt, seed, width, height, model_name, base_repo_id_or_path, pbar = gr.Progress(track_tqdm=True)):

train.py

@@ -1,3 +1,5 @@
+from torch.cuda.amp import autocast
+
 from StableDiffuser import StableDiffuser
 from finetuning import FineTunedModel
 import torch
@@ -8,20 +10,17 @@ 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
 
+    nsteps = 50
     diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=repo_id_or_path).to('cuda')
 
     memory_efficiency_wrapper = MemoryEfficiencyWrapper(diffuser=diffuser, use_amp=use_amp, use_xformers=use_xformers,
                                                         use_gradient_checkpointing=use_gradient_checkpointing )
     with memory_efficiency_wrapper:
-
         diffuser.train()
-
         finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
-
         if use_adamw8bit:
+            use print("using AdamW 8Bit optimizer")
             import bitsandbytes as bnb
             optimizer = bnb.optim.AdamW8bit(finetuner.parameters(),
                                             lr=lr,
@@ -30,13 +29,13 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                                             eps=1e-8
                                             )
         else:
+            print("using Adam optimizer")
             optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
         criteria = torch.nn.MSELoss()
 
         pbar = tqdm(range(iterations))
 
         with torch.no_grad():
-
             neutral_text_embeddings = diffuser.get_text_embeddings([''],n_imgs=1)
             positive_text_embeddings = diffuser.get_text_embeddings([prompt],n_imgs=1)
 
@@ -56,7 +55,7 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                 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:
+                with autocast(enabled=use_amp), finetuner:
                     latents_steps, _ = diffuser.diffusion(
                         latents,
                         positive_text_embeddings,
@@ -67,19 +66,21 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                     )
 
                 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 autocast(enabled=use_amp):
+                    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)
+                with autocast(enabled=use_amp):
+                    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 = criteria(e_n, e_0) works the best try 5000 epochs
+            loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents)))
             memory_efficiency_wrapper.step(optimizer, loss)
             optimizer.step()