Update app.py

app.py

@@ -14,6 +14,7 @@ from huggingface_hub import hf_hub_download
 # Ensure 'checkpoint' directory exists
 os.makedirs("checkpoints", exist_ok=True)
 
+# Download LoRA weights
 hf_hub_download(
     repo_id="wenqsun/DimensionX",
     filename="orbit_left_lora_weights.safetensors",
@@ -26,93 +27,78 @@ hf_hub_download(
     local_dir="checkpoints"
 )
 
+# Load models in the global scope
 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)
+transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16).to("cpu")
+text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16).to("cpu")
+vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16).to("cpu")
 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)
 
 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")
-
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+    gc.collect()
 
 def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True)):
     
-    pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
+    
     lora_path = "checkpoints/"
-    adapter_name = None
     if orbit_type == "Left":
         weight_name = "orbit_left_lora_weights.safetensors"
     elif orbit_type == "Up":
         weight_name = "orbit_up_lora_weights.safetensors"
     lora_rank = 256
 
+    pipe.unload_lora_weights()
+
     # Generate a timestamp for adapter_name
     adapter_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
 
+    # Load LoRA weights on CPU, move to GPU afterward
     pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=f"{adapter_timestamp}")
     pipe.fuse_lora(lora_scale=1 / lora_rank)
+    
+    # Move the pipeline to GPU for inference
     pipe.to("cuda")
-
-
+    
+    # Set the inference prompt
     prompt = f"{prompt}. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
     image = load_image(image_path)
     seed = random.randint(0, 2**8 - 1)
     
+    
     video = pipe(
         image, 
         prompt, 
-        num_inference_steps=50,  # NOT Changed
-        guidance_scale=7.0,  # NOT Changed
+        num_inference_steps=50,
+        guidance_scale=7.0,
         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
+    # Generate and save output video
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     export_to_video(video.frames[0], f"output_{timestamp}.mp4", fps=8)
+
+    # Move objects to CPU and clear GPU memory immediately after inference
+    find_and_move_object_to_cpu()
+    clear_gpu()
     
     return f"output_{timestamp}.mp4"
 
+# Set up Gradio UI
 with gr.Blocks(analytics_enabled=False) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# DimensionX")