Switch to text-only causal LM loading

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 import spaces
 import torch
 from transformers import (
-    AutoModelForImageTextToText,
+    AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
     TextIteratorStreamer,
@@ -29,6 +29,7 @@ PLACEHOLDER = (
 MAX_INPUT_TOKENS = 16384
 DEFAULT_MAX_NEW_TOKENS = 4096
 MAX_NEW_TOKENS = 8192
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -40,21 +41,27 @@ BNB_CONFIG = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-model = AutoModelForImageTextToText.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     trust_remote_code=True,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN,
+    device_map={"": 0},
+    dtype=torch.bfloat16,
     quantization_config=BNB_CONFIG,
     attn_implementation="sdpa",
+    low_cpu_mem_usage=True,
 )
 model.eval()
 
 
+def model_input_device():
+    return next(model.parameters()).device
+
+
 def estimate_duration(
     message,
     history,
@@ -113,7 +120,7 @@ def stream_chat(
         return_tensors="pt",
         truncation=True,
         max_length=MAX_INPUT_TOKENS,
-    ).to(model.device)
+    ).to(model_input_device())
 
     streamer = TextIteratorStreamer(
         tokenizer,