Update app.py

app.py

@@ -13,52 +13,59 @@ from transformers import AutoModel, AutoTokenizer
 from modelscope.hub.snapshot_download import snapshot_download
 from PIL import Image
 from decord import VideoReader, cpu
-import os
+
 import gc
 import io
 import tempfile
 from ultralytics import YOLO
 import numpy as np
 import cv2
+import os
+os.system("nvidia-smi")
 
 # Load YOLOv11 model (update the path as needed)
 YOLO_MODEL = YOLO('best_yolov11.pt')
 
-# Check if CUDA is available
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+model_dir = snapshot_download('iic/mPLUG-Owl3-7B-240728', cache_dir='./')
+os.system('ls')
+# README, How to run demo on different devices
 
-# Initialize GPU if available (already done by debug() above)
-if DEVICE == "cuda":
-    def init_debug():
-        torch.randn(10).cuda()
-    init_debug()
+# For Nvidia GPUs.
+# python web_demo_2.6.py --device cuda
 
-# File type validation
-IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
-VIDEO_EXTENSIONS = {'.mp4', '.mkv', '.mov', '.avi', '.flv', '.wmv', '.webm', '.m4v'}
+# For Mac with MPS (Apple silicon or AMD GPUs).
+# PYTORCH_ENABLE_MPS_FALLBACK=1 python web_demo_2.6.py --device mps
 
-def get_file_extension(filename):
-    return os.path.splitext(filename)[1].lower()
+# Argparser
+parser = argparse.ArgumentParser(description='demo')
+parser.add_argument('--device', type=str, default='cuda', help='cuda or mps')
+parser.add_argument("--host", type=str, default="0.0.0.0")
+parser.add_argument("--port", type=int)
+args = parser.parse_args()
+device = args.device
+assert device in ['cuda', 'mps']
 
-def is_image(filename):
-    return get_file_extension(filename) in IMAGE_EXTENSIONS
+# Load model
+model_path = './iic/mPLUG-Owl3-7B-240728'
+if 'int4' in model_path:
+    if device == 'mps':
+        print('Error: running int4 model with bitsandbytes on Mac is not supported right now.')
+        exit()
+    model = AutoModel.from_pretrained(model_path, attn_implementation='sdpa', trust_remote_code=True)
+else:
+    model = AutoModel.from_pretrained(model_path, attn_implementation='sdpa', trust_remote_code=True, torch_dtype=torch.bfloat16)
+    model = model.to(device=device)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model.eval()
 
-def is_video(filename):
-    return get_file_extension(filename) in VIDEO_EXTENSIONS
 
-# Model configuration
-MODEL_NAME = 'iic/mPLUG-Owl3-7B-240728'
-MODEL_CACHE_DIR = os.getenv('TRANSFORMERS_CACHE', './models')
-os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
-# Download and cache the model
-try:
-    model_path = snapshot_download(MODEL_NAME, cache_dir=MODEL_CACHE_DIR)
-except Exception as e:
-    print(f"Error downloading model: {str(e)}")
-    model_path = os.path.join(MODEL_CACHE_DIR, MODEL_NAME)
 
-MAX_NUM_FRAMES = 32
+ERROR_MSG = "Error, please retry"
+model_name = 'mPLUG-Owl3'
+MAX_NUM_FRAMES = 64
+IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
+VIDEO_EXTENSIONS = {'.mp4', '.mkv', '.mov', '.avi', '.flv', '.wmv', '.webm', '.m4v'}
 
 def load_model_and_tokenizer():
     """Load a fresh instance of the model and tokenizer"""