waleko commited on
Commit
77f8033
1 Parent(s): cd9a590

try fix zerogpu

Browse files
Files changed (1) hide show
  1. webui.py +2 -1
webui.py CHANGED
@@ -19,6 +19,7 @@ from transformers import TextIteratorStreamer, pipeline, ImageToTextPipeline, Au
19
  import os
20
  from pix2tex.cli import LatexOCR
21
  from munch import Munch
 
22
 
23
 
24
  from infer import TikzDocument, TikzGenerator
@@ -109,7 +110,7 @@ def inference(
109
  )
110
 
111
  thread = ThreadPool(processes=1)
112
- async_result = thread.apply_async(generate, kwds=dict(image=image, streamer=streamer))
113
  generated_text = ""
114
  for new_text in streamer:
115
  generated_text += new_text
 
19
  import os
20
  from pix2tex.cli import LatexOCR
21
  from munch import Munch
22
+ import spaces
23
 
24
 
25
  from infer import TikzDocument, TikzGenerator
 
110
  )
111
 
112
  thread = ThreadPool(processes=1)
113
+ async_result = thread.apply_async(spaces.GPU(generate), kwds=dict(image=image, streamer=streamer))
114
  generated_text = ""
115
  for new_text in streamer:
116
  generated_text += new_text