Use checkpoint-compatible model runtime

app.py

@@ -10,7 +10,7 @@ import socket
 import time
 import atexit
 import torch
-from transformers import AutoModel, AutoProcessor
+from transformers import AutoProcessor, MusicFlamingoForConditionalGeneration
 
 PROXY_URL = None
 _tunnel_proc = None
@@ -289,13 +289,11 @@ EXAMPLE_YOUTUBE_PROMPTS = [
 ]
 processor = AutoProcessor.from_pretrained(MODEL_ID)
 DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-DEVICE_MAP = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModel.from_pretrained(
+model = MusicFlamingoForConditionalGeneration.from_pretrained(
     MODEL_ID,
-    torch_dtype=DTYPE,
-    device_map=DEVICE_MAP,
+    dtype=DTYPE,
+    device_map="auto",
 ).eval()
-# model = AutoModel.from_pretrained(MODEL_ID, device_map="auto").eval()
 
 _youtube_cache = {}
 

requirements.txt

@@ -1,4 +1,4 @@
-transformers
+transformers==5.6.0
 accelerate
 torch
 torchaudio