More refactoring

app.py

@@ -1,11 +1,9 @@
-from pathlib import Path
-
 import gradio as gr
 import torch
 from finetuning import FineTunedModel
 from StableDiffuser import StableDiffuser
 from tqdm import tqdm
-
+from train import train
 
 model_map = {
     'Car' : 'models/car.pt',
@@ -18,41 +16,16 @@ class Demo:
     def __init__(self) -> None:
 
         self.training = False
-        self.generating = False
-        self.nsteps = 50
 
-        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda')
-        self.finetuner = None
-        
+        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
 
         with gr.Blocks() as demo:
             self.layout()
-            self.switch_model(self.model_dropdown.value)
-
-            self.finetuner = self.finetuner.eval().half()
-            self.diffuser = self.diffuser.eval().half()
-
             demo.queue(concurrency_count=2).launch()
 
-    def disable(self):
-        return [gr.update(interactive=False), gr.update(interactive=False)]
-  
-    def switch_model(self, model_name):
-
-        if not model_name:
-            return
-
-        model_path = model_map[model_name]
-        
-        checkpoint = torch.load(model_path)
-
-        self.finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint)
-
-        torch.cuda.empty_cache()
 
     def layout(self):
 
-       
         with gr.Row():
             
                 
@@ -149,25 +122,24 @@ class Demo:
 
                     with gr.Column(scale=1):
 
+                        self.train_status = gr.Button(value='', variant='primary', label='Status', interactive=False)
+
                         self.train_button = gr.Button(
                             value="Train",
                         )
 
                         self.download = gr.Files()
 
-        self.model_dropdown.change(self.switch_model, inputs=[self.model_dropdown])
         self.infr_button.click(self.inference, inputs = [
             self.prompt_input_infr,
-            self.seed_infr
+            self.seed_infr,
+            self.model_dropdown
             ],
             outputs=[
                 self.image_new,
                 self.image_orig
             ]
         )
-        self.train_button.click(self.disable, 
-            outputs=[self.train_button, self.infr_button]
-        )
         self.train_button.click(self.train, inputs = [
             self.prompt_input,
             self.train_method_input, 
@@ -175,21 +147,13 @@ class Demo:
             self.iterations_input,
             self.lr_input
         ],
-        outputs=[self.train_button, self.infr_button, self.download, self.model_dropdown]
+        outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
         )
 
     def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
 
         if self.training:
-            return [None, None, None]
-        else:
-            self.training = True
-
-        del self.finetuner
-
-        torch.cuda.empty_cache()
-
-        self.diffuser = self.diffuser.train().float()
+            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':
 
@@ -206,82 +170,35 @@ class Demo:
             modules = ".*attn1$"
             frozen = []
 
-        finetuner = FineTunedModel(self.diffuser, modules, frozen_modules=frozen)
-
-        optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
-        criteria = torch.nn.MSELoss()
-
-        pbar = tqdm(range(iterations))
+        randn = torch.randint(1, 10000000, (1,)).item()
 
-        with torch.no_grad():
+        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}.pt"
 
-            neutral_text_embeddings = self.diffuser.get_text_embeddings([''],n_imgs=1)
-            positive_text_embeddings = self.diffuser.get_text_embeddings([prompt],n_imgs=1)
+        self.training = True
 
-        for i in pbar:
-            
-            with torch.no_grad():
-
-                self.diffuser.set_scheduler_timesteps(self.nsteps)
-
-                optimizer.zero_grad()
-
-                iteration = torch.randint(1, self.nsteps - 1, (1,)).item()
-
-                latents = self.diffuser.get_initial_latents(1, 512, 1)
-
-                with finetuner:
+        train(prompt, modules, frozen, iterations, neg_guidance, lr, save_path)
 
-                    latents_steps, _ = self.diffuser.diffusion(
-                        latents,
-                        positive_text_embeddings,
-                        start_iteration=0,
-                        end_iteration=iteration,
-                        guidance_scale=3, 
-                        show_progress=False
-                    )
+        self.training = False
 
-                self.diffuser.set_scheduler_timesteps(1000)
+        torch.cuda.empty_cache()
 
-                iteration = int(iteration / self.nsteps * 1000)
-                
-                positive_latents = self.diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
-                neutral_latents = self.diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1)
+        model_map['Custom'] = save_path
 
-            with finetuner:
-                negative_latents = self.diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
 
-            positive_latents.requires_grad = False
-            neutral_latents.requires_grad = False
 
-            loss = criteria(negative_latents, neutral_latents - (neg_guidance*(positive_latents - neutral_latents))) #loss = criteria(e_n, e_0) works the best try 5000 epochs
-            loss.backward()
-            optimizer.step()
+    def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
 
-        ft_path = f"{prompt.lower().replace(' ', '')}.pt"
-        torch.save(finetuner.state_dict(), ft_path)
+        self.diffuser._seed = seed or 42
 
-        self.finetuner = finetuner.eval().half()
+        model_path = model_map[model_name]
+        
+        checkpoint = torch.load(model_path)
 
-        self.diffuser = self.diffuser.eval().half()
+        self.finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
 
         torch.cuda.empty_cache()
 
-        self.training = False      
-
-        model_map['Custom'] = ft_path
-
-        return [gr.update(interactive=True), gr.update(interactive=True), ft_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
-
-
-    def inference(self, prompt, seed, pbar = gr.Progress(track_tqdm=True)):
-        if self.generating:
-            return [None, None]
-        else:
-            self.generating = True
-
-        self.diffuser._seed = seed or 42
-
         images = self.diffuser(
             prompt,
             n_steps=50,
@@ -302,8 +219,6 @@ class Demo:
 
         edited_image = images[0][0]
 
-        self.generating = False
-
         torch.cuda.empty_cache()
 
         return edited_image, orig_image

train.py

@@ -0,0 +1,84 @@
+from StableDiffuser import StableDiffuser
+from finetuning import FineTunedModel
+import torch
+from tqdm import tqdm
+
+def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path):
+  
+    nsteps = 50
+
+    diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
+    diffuser.train()
+
+    finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
+
+    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)
+
+    losses = []
+
+    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, 512, 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()
+        losses.append(loss.item())
+        optimizer.step()
+
+    torch.save(finetuner.state_dict(), save_path)
+
+if __name__ == '__main__':
+
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--prompt', required=True)
+    parser.add_argument('--modules', required=True)
+    parser.add_argument('--freeze_modules', nargs='+', required=True)
+    parser.add_argument('--save_path', required=True)
+    parser.add_argument('--iterations', type=int, required=True)
+    parser.add_argument('--lr', type=float, required=True)
+    parser.add_argument('--negative_guidance', type=float, required=True)
+
+    train(**vars(parser.parse_args()))