Update app.py

app.py

@@ -89,7 +89,7 @@ def reset_components():
 
 
 @spaces.GPU
-def main(video, prompt, role, temperature, max_new_tokens, model, processor, device):
+def main(video, prompt, role, temperature, max_new_tokens, model, processor):
     history = []
 
     if not video:
@@ -159,7 +159,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, dev
 
         images, videos = process_vision_info(messages)
         data = processor(text=[text], images=images, videos=videos, return_tensors='pt')
-        data = data.to(device)
+        data = data.to(next(model.parameters()).device)
 
         model.base_model.disable_adapter_layers()
         model.base_model.enable_adapter_layers()
@@ -258,7 +258,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, dev
         text = processor.apply_chat_template(messages, add_generation_prompt=True)
         images, videos = process_vision_info(messages)
         data = processor(text=[text], images=images, videos=videos, return_tensors='pt')
-        data = data.to(device)
+        data = data.to(next(model.parameters()).device)
 
         model.base_model.disable_adapter_layers()
         model.base_model.enable_adapter_layers()
@@ -395,7 +395,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, dev
             data['attention_mask'] = torch.ones_like(data['input_ids'])
             # ===========================================
 
-            data = data.to(device)
+            data = data.to(next(model.parameters()).device)
 
             model.base_model.disable_adapter_layers()
             model.base_model.enable_adapter_layers()
@@ -482,7 +482,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, dev
         text = processor.apply_chat_template(messages, add_generation_prompt=True)
         images, videos = process_vision_info(messages)
         data = processor(text=[text], images=images, videos=videos, return_tensors='pt')
-        data = data.to(device)
+        data = data.to(next(model.parameters()).device)
 
         with model.disable_adapter():
             output_ids = model.generate(
@@ -539,9 +539,7 @@ if __name__ == '__main__':
     print('Initializing role *verifier*')
     model.load_adapter(nncore.join(MODEL, 'verifier'), adapter_name='verifier')
 
-    device = next(model.parameters()).device
-
-    main = partial(main, model=model, processor=processor, device=device)
+    main = partial(main, model=model, processor=processor)
 
     path = os.path.dirname(os.path.realpath(__file__))