blehhhhh

app.py

@@ -16,9 +16,9 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 
 @spaces.GPU
 def tag_image(image, max_new_tokens=128):
-    inputs = processor(images=image, text="<image>tag en", return_tensors="pt")
+    inputs = processor(images=image, text="<image>tag en", return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer=processor.tokenizer, skip_prompt=True)
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens, use_cache=True, cache_implementation="hybrid")
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     
     text = ""
@@ -26,8 +26,6 @@ def tag_image(image, max_new_tokens=128):
     for new_text in streamer:
         text += new_text
         yield text
-    
-    thread.join()
 
     return text