Update app.py

app.py

@@ -13,17 +13,21 @@ DEPLOY_MODELS = {}
 IMAGE_TOKEN = "<image>"
 
 # Fetch model
-def fetch_model(model_name: str, dtype=torch.bfloat16):
+def fetch_model(model_name: str, dtype=None):
     global DEPLOY_MODELS
     if model_name not in DEPLOY_MODELS:
-        logger.info(f"Loading {model_name}...")
+        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}...")
         model_info = load_model(model_name, dtype=dtype)
         tokenizer, model, vl_chat_processor = model_info
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         model = model.to(device)
         DEPLOY_MODELS[model_name] = (tokenizer, model, vl_chat_processor)
-        logger.info(f"Loaded {model_name} on {device}")
+        logger.info(f"Loaded {model_name} successfully.")
     return DEPLOY_MODELS[model_name]
+
     
 # Generate prompt with history
 def generate_prompt_with_history(text, images, history, vl_chat_processor, tokenizer, max_length=2048):