v1

app.py

@@ -14,19 +14,9 @@ from transformers import TextIteratorStreamer
 from torchvision.transforms.functional import pil_to_tensor
 
 # loading meteor model
-mmamba = load_mmamba('BK-Lee/Meteor-Mamba').to('cuda')
+mmamba = load_mmamba('BK-Lee/Meteor-Mamba').cuda()
 meteor, tok_meteor = load_meteor('BK-Lee/Meteor-MLM', bits=16)
 
-# param
-for param in mmamba.parameters():
-    param = param.to('cuda')
-for param in meteor.parameters():
-    param = param.to('cuda')
-
-# device
-device = torch.cuda.current_device()
-
-
 # freeze model
 freeze_model(mmamba)
 freeze_model(meteor)
@@ -36,7 +26,16 @@ previous_length = 0
 
 @spaces.GPU
 def threading_function(inputs, image_token_number, streamer):
-    print(f'----------------------------Device: {device}----------------------------')
+
+    # device
+    device = torch.cuda.current_device()
+
+    # param
+    for param in mmamba.parameters():
+        param.data = param.to(device)
+    for param in meteor.parameters():
+        param.data = param.to(device)
+
     # Meteor Mamba
     mmamba_inputs = mmamba.eval_process(inputs=inputs, tokenizer=tok_meteor, device=device, img_token_number=image_token_number)
     if 'image' in mmamba_inputs.keys():

meteor/arch/modeling_internlm2.py

@@ -277,9 +277,6 @@ def rotate_half(x):
 # Copied from transformers.model.llama.modeling_llama.apply_rotary_pos_emb
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors."""
-    print(f'------------------------------cos: {cos.device}------------------------------')
-    print(f'------------------------------position_ids: {position_ids.device}------------------------------')
-    print(f'------------------------------unsqueeze_dim: {unsqueeze_dim.device}------------------------------')
     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)