Update demo/infer.py

demo/infer.py

@@ -32,12 +32,13 @@ class LiveCCDemoInfer:
     streaming_time_interval = streaming_fps_frames / fps
     frame_time_interval = 1 / fps
 
-    def __init__(self, model_path: str = None, device_id: int = 0):
+    def __init__(self, model_path: str = None, device: str = 'cpu'):
         self.model = Qwen2VLForConditionalGeneration.from_pretrained(
             model_path, torch_dtype="auto", 
             # device_map=f'cuda:{device_id}', 
             # attn_implementation='flash_attention_2'
         )
+        self.model.to(device)
         self.processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
         self.streaming_eos_token_id = self.processor.tokenizer(' ...').input_ids[-1]
         self.model.prepare_inputs_for_generation = functools.partial(prepare_multiturn_multimodal_inputs_for_generation, self.model)