Update app.py

app.py

@@ -16,23 +16,35 @@ from qwen_omni_utils import process_mm_info
 from argparse import ArgumentParser
 
 def _load_model_processor(args):
+    import torch
+
     if args.cpu_only:
         device_map = 'cpu'
+        max_memory = {0: "2GB"}  # Limit memory usage when running on CPU
     else:
         device_map = 'auto'
+        max_memory = {i: "20GB" for i in range(torch.cuda.device_count())}  # Adjust as needed
 
     # Check if flash-attn2 flag is enabled and load model accordingly
     if args.flash_attn2:
-        model = Qwen2_5OmniModel.from_pretrained(args.checkpoint_path,
-                                                    torch_dtype='auto',
-                                                    attn_implementation='flash_attention_2',
-                                                    device_map=device_map)
+        model = Qwen2_5OmniModel.from_pretrained(
+            args.checkpoint_path,
+            torch_dtype='auto',
+            attn_implementation='flash_attention_2',
+            device_map=device_map,
+            max_memory=max_memory
+        )
     else:
-        model = Qwen2_5OmniModel.from_pretrained(args.checkpoint_path, device_map=device_map)
+        model = Qwen2_5OmniModel.from_pretrained(
+            args.checkpoint_path,
+            device_map=device_map,
+            max_memory=max_memory
+        )
 
     processor = Qwen2_5OmniProcessor.from_pretrained(args.checkpoint_path)
     return model, processor
 
+
 def _launch_demo(args, model, processor):
     # Voice settings
     VOICE_LIST = ['Chelsie', 'Ethan']