Update app.py

app.py

@@ -12,13 +12,13 @@ from huggingface_hub import snapshot_download
 MODEL_BASE_SINGLE = snapshot_download(repo_id="nvidia/audio-flamingo-3")
 MODEL_BASE_THINK = os.path.join(MODEL_BASE_SINGLE, 'stage35')
 
-model_single = llava.load(MODEL_BASE_SINGLE, model_base=None, devices=[0])
+model_single = llava.load(MODEL_BASE_SINGLE, model_base=None, devices=[0,1])
 generation_config_single = model_single.default_generation_config
 
 model_think = PeftModel.from_pretrained(
             model_single,
             MODEL_BASE_THINK,
-            device_map="auto",
+            device_map="cuda:2",
             torch_dtype=torch.float16,
         )
 
@@ -26,7 +26,7 @@ model_think = PeftModel.from_pretrained(
 # MULTI-TURN MODEL SETUP
 # ---------------------------------
 MODEL_BASE_MULTI = snapshot_download(repo_id="nvidia/audio-flamingo-3-chat")
-model_multi = llava.load(MODEL_BASE_MULTI, model_base=None, devices=[0])
+model_multi = llava.load(MODEL_BASE_MULTI, model_base=None, devices=[3])
 generation_config_multi = model_multi.default_generation_config