Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 import spaces                # <-- NEW: import spaces for ZeroGPU
 import gradio as gr
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForVision2Seq, AutoConfig
+from transformers import AutoProcessor, AutoModelForVision2Seq, AutoConfig, AutoModelForCausalLM
 from peft import PeftModel
 import transformers, accelerate, torch, numpy as np
 
@@ -27,7 +27,7 @@ processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
 
 cfg = AutoConfig.from_pretrained(BASE_MODEL, trust_remote_code=True)
 # 3. Load base model on CPU; stream shards to save RAM
-base = AutoModelForVision2Seq.from_pretrained(
+base = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     config=cfg,
     torch_dtype=dtype,