Dec GPU usage

app.py

@@ -1,15 +1,15 @@
-from gradio import ChatInterface, Textbox, Slider
-from spaces import GPU
 from threading import Thread
-from torch import bfloat16
-from transformers import Qwen2VLForConditionalGeneration, Qwen2VLProcessor, TextIteratorStreamer, AutoProcessor
+from spaces import GPU
+from gradio import ChatInterface, Textbox, Slider
+
+from transformers import Qwen2VLForConditionalGeneration, Qwen2VLProcessor, TextIteratorStreamer, AutoProcessor, BatchFeature
 from qwen_vl_utils import process_vision_info
 
 model_path = "Pectics/Softie-VL-7B-250123"
 
 model = Qwen2VLForConditionalGeneration.from_pretrained(
     model_path,
-    torch_dtype=bfloat16,
+    torch_dtype="auto",
     attn_implementation="flash_attention_2",
     device_map="auto",
 )
@@ -18,53 +18,45 @@ max_pixels = 1280 * 28 * 28
 processor: Qwen2VLProcessor = AutoProcessor.from_pretrained(model_path, min_pixels=min_pixels, max_pixels=max_pixels)
 
 @GPU
-def infer(
-    messages,
+def infer(inputs: BatchFeature, **kwargs) -> None:
+    inputs = inputs.to("cuda")
+    model.generate(**kwargs)
+
+def respond(
+    message,
+    history,
+    system_message,
     max_tokens,
     temperature,
     top_p,
 ):
+    messages = [{"role": "system", "content": system_message}]
+    for m in history:
+        messages.append({"role": m["role"], "content": m["content"]})
+    messages.append({"role": "user", "content": message})
     text_inputs = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(
-        text=[text_inputs],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    ).to("cuda")
-    
+        text = [text_inputs],
+        images = image_inputs,
+        videos = video_inputs,
+        padding = True,
+        return_tensors = "pt",
+    )
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
     kwargs = dict(
-        **inputs,
+        inputs=inputs,
         streamer=streamer,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
     )
-    thread = Thread(target=model.generate, kwargs=kwargs)
-    thread.start()
-    
+    Thread(target=infer, kwargs=kwargs).start()
     response = ""
     for token in streamer:
         response += token
         yield response
 
-def respond(
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-    for m in history:
-        messages.append({"role": m["role"], "content": m["content"]})
-    messages.append({"role": "user", "content": message})
-    for response in infer(messages, max_tokens, temperature, top_p):
-        yield response
-
 app = ChatInterface(
     respond,
     type="messages",