v1

app.py

@@ -8,13 +8,17 @@ from PIL import Image
 from utils.utils import *
 from threading import Thread
 import torch.nn.functional as F
+from accelerate import Accelerator
 from meteor.load_mmamba import load_mmamba
 from meteor.load_meteor import load_meteor
 from transformers import TextIteratorStreamer
 from torchvision.transforms.functional import pil_to_tensor
 
+# accel
+accel = Accelerator()
+
 # loading meteor model
-mmamba = load_mmamba('BK-Lee/Meteor-Mamba').cuda()
+mmamba = load_mmamba('BK-Lee/Meteor-Mamba')
 meteor, tok_meteor = load_meteor('BK-Lee/Meteor-MLM', bits=16)
 
 # freeze model
@@ -24,7 +28,6 @@ freeze_model(meteor)
 # previous length
 previous_length = 0
 
-@spaces.GPU
 def threading_function(inputs, image_token_number, streamer, device):
 
     # Meteor Mamba
@@ -49,24 +52,14 @@ def threading_function(inputs, image_token_number, streamer, device):
     generation_kwargs.update({'use_cache': True})
     return meteor.generate(**generation_kwargs)
 
-def add_message(history, message):
-    for x in message["files"]:
-        history.append(((x,), None))
-    if message["text"] is not None:
-        history.append((message["text"], None))
-    return history, gr.MultimodalTextbox(value=None, interactive=False)
-
 @spaces.GPU
 def bot_streaming(message, history):
 
-    # device
-    device = torch.cuda.current_device()
-
     # param
     for param in mmamba.parameters():
-        param.data = param.to(device)
+        param.data = param.to(accel.device)
     for param in meteor.parameters():
-        param.data = param.to(device)
+        param.data = param.to(accel.device)
 
     # prompt type -> input prompt
     image_token_number = int((490/14)**2)
@@ -83,7 +76,7 @@ def bot_streaming(message, history):
         streamer = TextIteratorStreamer(tok_meteor, skip_special_tokens=True)
 
         # Threading generation
-        thread = Thread(target=threading_function, kwargs=dict(inputs=inputs, image_token_number=image_token_number, streamer=streamer, device=device))
+        thread = Thread(target=threading_function, kwargs=dict(inputs=inputs, image_token_number=image_token_number, streamer=streamer, device=accel.device))
         thread.start()
 
         # generated text

meteor/arch/modeling_internlm2.py

@@ -277,8 +277,8 @@ 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."""
-    cos = cos.to(position_ids.device)[position_ids].unsqueeze(unsqueeze_dim)
-    sin = sin.to(position_ids.device)[position_ids].unsqueeze(unsqueeze_dim)
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed

requirements.txt

@@ -13,4 +13,5 @@ timm
 shortuuid
 matplotlib
 gradio
-spaces
+spaces
+accelerate