Update app.py

app.py

@@ -13,19 +13,22 @@ DEPLOY_MODELS = {}
 IMAGE_TOKEN = "<image>"
 
 # Fetch model
-def fetch_model(model_name: str, dtype=None):
+def fetch_model(model_name: str, dtype=torch.bfloat16):
     global DEPLOY_MODELS
     if model_name not in DEPLOY_MODELS:
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        # Use bfloat16 only if using GPU
-        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-        logger.info(f"Loading {model_name} on {device} with dtype={dtype}...")
+        logger.info(f"Loading {model_name}...")
         model_info = load_model(model_name, dtype=dtype)
         tokenizer, model, vl_chat_processor = model_info
-        model = model.to(device)
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        try:
+            model = model.to(device)
+        except RuntimeError as e:
+            logger.warning(f"Could not move model to {device}: {e}")
+            device = torch.device('cpu')
+            model = model.to(device)
+            logger.warning("Model fallback to CPU. Inference might be slow.")
         DEPLOY_MODELS[model_name] = (tokenizer, model, vl_chat_processor)
-        logger.info(f"Loaded {model_name} successfully.")
+        logger.info(f"Loaded {model_name} on {device}")
     return DEPLOY_MODELS[model_name]