Memory saving

train.py

@@ -10,6 +10,9 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
     diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
     diffuser.train()
 
+
+
+
     finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
 
     optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
@@ -22,7 +25,11 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
         neutral_text_embeddings = diffuser.get_text_embeddings([''],n_imgs=1)
         positive_text_embeddings = diffuser.get_text_embeddings([prompt],n_imgs=1)
 
-    losses = []
+    del diffuser.vae
+    del diffuser.text_encoder
+    del diffuser.tokenizer
+
+    torch.cuda.empty_cache()
 
     for i in pbar:
         
@@ -61,8 +68,11 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
         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
+        
+        del negative_latents, neutral_latents, positive_latents, latents_steps, latents
+        torch.cuda.empty_cache()
+
         loss.backward()
-        losses.append(loss.item())
         optimizer.step()
 
     torch.save(finetuner.state_dict(), save_path)