add extra gc

hydra_node/models.py

@@ -42,10 +42,15 @@ def pil_upscale(image, scale=1):
 def fix_batch(tensor, bs):
     return torch.stack([tensor.squeeze(0)]*bs, dim=0)
 
+def memory_usage_display():
+    return f"Allocated {torch.cuda.memory_allocated() / 1024 / 1024}M, Reserved {torch.cuda.memory_reserved() / 1024 / 1024}M"
+
 def torch_gc():
     if torch.cuda.is_available():
+        #print(f"Before clean: {memory_usage_display()}")
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
+        #print(f"After clean: {memory_usage_display()}")
 
 # make uc and prompt shapes match via padding for long prompts
 # finetune
@@ -482,6 +487,7 @@ class StableDiffusionModel(nn.Module):
         #set hypernetwork to none after generation
         CrossAttention.set_hypernetwork(None)
 
+        torch_gc()
         return images
 
     @torch.no_grad()
@@ -579,6 +585,7 @@ class StableDiffusionModel(nn.Module):
             torch.seed()
             np.random.seed()
 
+        torch_gc()
         return images
 
     @torch.no_grad()