code cleanup

server.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoProcessor
+from transformers import AutoProcessor
 import os
 import torch
 
@@ -11,40 +11,17 @@ CACHE_DIR = (
     or "/data/huggingface"
 )
 
-# Inspect config and load appropriate stack
-config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
+from transformers import Qwen2VLForConditionalGeneration  # type: ignore
 
-is_qwen2_vl = getattr(config, "model_type", None) == "qwen2_vl" or (
-    config.__class__.__name__.lower().startswith("qwen2vl")
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+    cache_dir=CACHE_DIR,
 )
 
-if is_qwen2_vl:
-    try:
-        from transformers import Qwen2VLForConditionalGeneration  # type: ignore
-    except Exception as e:
-        raise RuntimeError(
-            "Transformers version does not support Qwen2-VL. Please upgrade transformers to >=4.43."
-        ) from e
-    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
-    model = Qwen2VLForConditionalGeneration.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True,
-        cache_dir=CACHE_DIR,
-    )
-    _use_processor = True
-else:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",  # automatically shards across GPUs
-        trust_remote_code=True,
-        cache_dir=CACHE_DIR
-    )
-    _use_processor = False
-
 app = FastAPI()
 
 # OpenAI-style request schema
@@ -61,14 +38,9 @@ class ChatCompletionRequest(BaseModel):
 async def chat_completions(req: ChatCompletionRequest):
     # Concatenate messages into one prompt
     prompt = "\n".join([m.content for m in req.messages])
-    if _use_processor:
-        inputs = processor(text=prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
-        text = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
-    else:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
-        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = processor(text=prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=req.max_tokens)
+    text = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     return {
         "id": "chatcmpl-uground72b",