complete memory management

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import os
 import torch
+import gc
 from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from diffusers.utils import export_to_video, load_image
 from transformers import T5EncoderModel, T5Tokenizer
@@ -25,17 +26,54 @@ hf_hub_download(
     local_dir="checkpoints"
 )
 
-
 model_id = "THUDM/CogVideoX-5b-I2V"
 transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
 text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16)
 vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
 tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
 pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
-lora_path = "your lora path"
-lora_rank = 256
+
+def find_and_move_object_to_cpu():
+    for obj in gc.get_objects():
+        try:
+            # Check if the object is a PyTorch model
+            if isinstance(obj, torch.nn.Module):
+                # Check if any parameter of the model is on CUDA
+                if any(param.is_cuda for param in obj.parameters()):
+                    print(f"Found PyTorch model on CUDA: {type(obj).__name__}")
+                    # Move the model to CPU
+                    obj.to('cpu')
+                    print(f"Moved {type(obj).__name__} to CPU.")
+                    
+                # Optionally check if buffers are on CUDA
+                if any(buf.is_cuda for buf in obj.buffers()):
+                    print(f"Found buffer on CUDA in {type(obj).__name__}")
+                    obj.to('cpu')
+                    print(f"Moved buffers of {type(obj).__name__} to CPU.")
+
+        except Exception as e:
+            # Handle any exceptions if obj is not a torch model
+            pass
+
+
+def clear_gpu():
+    """Clear GPU memory by removing tensors, freeing cache, and moving data to CPU."""
+    # List memory usage before clearing
+    print(f"Memory allocated before clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
+    print(f"Memory reserved before clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
+
+    # Move any bound tensors back to CPU if needed
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()  # Ensure that all operations are completed
+        print("GPU memory cleared.")
+    
+    print(f"Memory allocated after clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
+    print(f"Memory reserved after clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
+
 
 def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True)):
+
     lora_path = "checkpoints/"
     adapter_name = None
     if orbit_type == "Left":
@@ -64,10 +102,14 @@ def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True))
         use_dynamic_cfg=True,
         generator=torch.Generator(device="cpu").manual_seed(seed)
     )
+
+    find_and_move_object_to_cpu()
+    clear_gpu()
     
     # Generate a timestamp for the output filename
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     export_to_video(video.frames[0], f"output_{timestamp}.mp4", fps=8)
+    
     return f"output_{timestamp}.mp4"
 
 with gr.Blocks() as demo: