Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
 import deepspeed
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Model name
 model_name = "OpenGVLab/InternVideo2_5_Chat_8B"
@@ -9,19 +9,27 @@ model_name = "OpenGVLab/InternVideo2_5_Chat_8B"
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# Load model efficiently
+# Enable DeepSpeed Inference (ZeRO-3)
+ds_engine = deepspeed.init_inference(
+    dtype=torch.float16,     # Use float16 for efficiency
+    replace_method="auto",   # Automatically replace ops for inference
+    replace_with_kernel_inject=True
+)
+
+# Load model with DeepSpeed
 model = AutoModelForCausalLM.from_pretrained(
     model_name, 
     trust_remote_code=True,
-    torch_dtype=torch.float16,  # Use float16 for lower memory usage
-    device_map="auto", # Automatically place model on available GPU
-    deepspeed={"stage": 3}      # Enable DeepSpeed ZeRO-3
-    
+    torch_dtype=torch.float16, 
+    device_map="auto"  # Auto place on GPU
 )
 
+# Apply DeepSpeed to model
+model = ds_engine.module(model)
+
 # Define inference function
 def chat_with_model(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Move inputs to GPU
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     output = model.generate(**inputs, max_length=200)
     return tokenizer.decode(output[0], skip_special_tokens=True)