Update model_worker.py
Browse files- model_worker.py +3 -1
model_worker.py
CHANGED
|
@@ -29,6 +29,8 @@ from transformers import TextIteratorStreamer
|
|
| 29 |
from threading import Thread
|
| 30 |
from icecream import ic
|
| 31 |
|
|
|
|
|
|
|
| 32 |
|
| 33 |
GB = 1 << 30
|
| 34 |
|
|
@@ -70,7 +72,7 @@ class ModelWorker:
|
|
| 70 |
|
| 71 |
self.is_multimodal = True
|
| 72 |
|
| 73 |
-
|
| 74 |
@torch.inference_mode()
|
| 75 |
def generate_stream(self, params):
|
| 76 |
tokenizer, model = self.tokenizer, self.model
|
|
|
|
| 29 |
from threading import Thread
|
| 30 |
from icecream import ic
|
| 31 |
|
| 32 |
+
import spaces # for use zero of huggingface
|
| 33 |
+
|
| 34 |
|
| 35 |
GB = 1 << 30
|
| 36 |
|
|
|
|
| 72 |
|
| 73 |
self.is_multimodal = True
|
| 74 |
|
| 75 |
+
@spaces.GPU
|
| 76 |
@torch.inference_mode()
|
| 77 |
def generate_stream(self, params):
|
| 78 |
tokenizer, model = self.tokenizer, self.model
|