Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import torch
-import deepspeed
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Model name
@@ -9,27 +8,23 @@ model_name = "OpenGVLab/InternVideo2_5_Chat_8B"
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# 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
-)
+# Detect device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load model with DeepSpeed
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
+    model_name,
     trust_remote_code=True,
-    torch_dtype=torch.float16, 
-    device_map="auto"  # Auto place on GPU
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,  # Use float16 on GPU, float32 on CPU
+    device_map="auto" if device == "cuda" else None  # Use GPU if available
 )
 
-# Apply DeepSpeed to model
-model = ds_engine.module(model)
+# Move model to device
+model.to(device)
 
 # Define inference function
 def chat_with_model(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     output = model.generate(**inputs, max_length=200)
     return tokenizer.decode(output[0], skip_special_tokens=True)