v1

app.py

@@ -25,16 +25,7 @@ freeze_model(meteor)
 previous_length = 0
 
 @spaces.GPU
-def threading_function(inputs, image_token_number, streamer):
-
-    # 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)
+def threading_function(inputs, image_token_number, streamer, device):
 
     # Meteor Mamba
     mmamba_inputs = mmamba.eval_process(inputs=inputs, tokenizer=tok_meteor, device=device, img_token_number=image_token_number)
@@ -68,6 +59,15 @@ def add_message(history, message):
 @spaces.GPU
 def bot_streaming(message, history):
 
+    # 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)
+
     # prompt type -> input prompt
     image_token_number = int((490/14)**2)
     if len(message['files']) != 0:
@@ -83,7 +83,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))
+        thread = Thread(target=threading_function, kwargs=dict(inputs=inputs, image_token_number=image_token_number, streamer=streamer, device=device))
         thread.start()
 
         # generated text