refactor: update model loading to use AutoModelForVision2Seq instead of AutoModel

app.py

@@ -9,7 +9,7 @@ import gradio as gr
 from gradio import Server
 from fastapi.responses import HTMLResponse
 import torch
-from transformers import AutoProcessor, AutoModel
+from transformers import AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
 from dotenv import load_dotenv
 
@@ -37,7 +37,7 @@ from gradio.data_classes import FileData
 logger.info("Loading model and processor...")
 model_id = "HiDream-ai/HiDream-O1-Image"
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModel.from_pretrained(
+model = AutoModelForVision2Seq.from_pretrained(
     model_id, 
     torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
     trust_remote_code=True