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):
+def main(video, prompt, role, temperature, max_new_tokens, model, processor, device):
     history = []
 
     if not video:
@@ -115,6 +115,8 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor):
     history.append({'role': 'user', 'content': prompt})
     yield history
 
+    model = model.to(device)
+
     duration = get_duration(video)
 
     # do grounding and answering by default
@@ -159,9 +161,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor):
 
         images, videos = process_vision_info(messages)
         data = processor(text=[text], images=images, videos=videos, return_tensors='pt')
-        data = data.to(next(model.parameters()).device)
-
-        print(next(model.parameters()).device)
+        data = data.to(device)
 
         model.base_model.disable_adapter_layers()
         model.base_model.enable_adapter_layers()
@@ -260,7 +260,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor):
         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(next(model.parameters()).device)
+        data = data.to(device)
 
         model.base_model.disable_adapter_layers()
         model.base_model.enable_adapter_layers()
@@ -397,7 +397,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor):
             data['attention_mask'] = torch.ones_like(data['input_ids'])
             # ===========================================
 
-            data = data.to(next(model.parameters()).device)
+            data = data.to(device)
 
             model.base_model.disable_adapter_layers()
             model.base_model.enable_adapter_layers()
@@ -484,7 +484,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor):
         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(next(model.parameters()).device)
+        data = data.to(device)
 
         with model.disable_adapter():
             output_ids = model.generate(
@@ -541,7 +541,9 @@ if __name__ == '__main__':
     print('Initializing role *verifier*')
     model.load_adapter(nncore.join(MODEL, 'verifier'), adapter_name='verifier')
 
-    main = partial(main, model=model, processor=processor)
+    device = torch.device('cuda')
+
+    main = partial(main, model=model, processor=processor, device=device)
 
     path = os.path.dirname(os.path.realpath(__file__))