nxphi47 commited on
Commit
d31d891
1 Parent(s): 2a3225c

Update multipurpose_chatbot/engines/sealmmm_engine.py

Browse files
multipurpose_chatbot/engines/sealmmm_engine.py CHANGED
@@ -4,6 +4,7 @@
4
  import os
5
  import numpy as np
6
  import argparse
 
7
  import torch
8
  import gradio as gr
9
  from typing import Any, Iterator
@@ -201,7 +202,8 @@ class SeaLMMMv0Engine(TransformersEngine):
201
  message_safety = safety_check_conversation_string(message)
202
  if message_safety is not None:
203
  raise gr.Error(message_safety)
204
-
 
205
  def generate_yield_string(self, prompt, temperature, max_tokens, stop_strings: Optional[Tuple[str]] = None, **kwargs):
206
  from transformers.generation.utils import GenerationConfig
207
  from PIL import Image
 
4
  import os
5
  import numpy as np
6
  import argparse
7
+ import spaces
8
  import torch
9
  import gradio as gr
10
  from typing import Any, Iterator
 
202
  message_safety = safety_check_conversation_string(message)
203
  if message_safety is not None:
204
  raise gr.Error(message_safety)
205
+
206
+ @spaces.GPU
207
  def generate_yield_string(self, prompt, temperature, max_tokens, stop_strings: Optional[Tuple[str]] = None, **kwargs):
208
  from transformers.generation.utils import GenerationConfig
209
  from PIL import Image