KingNish commited on
Commit
713f201
1 Parent(s): e82c7af

Streaming output

Browse files
Files changed (1) hide show
  1. app.py +85 -25
app.py CHANGED
@@ -1,34 +1,81 @@
1
  import gradio as gr
2
  import spaces
3
- from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
4
  from qwen_vl_utils import process_vision_info
5
  import torch
6
  from PIL import Image
7
  import subprocess
8
  import numpy as np
9
  import os
10
-
11
- # Install flash-attn
12
- subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
13
 
14
  # Model and Processor Loading (Done once at startup)
15
  MODEL_ID = "Qwen/Qwen2-VL-7B-Instruct"
16
- model = Qwen2VLForConditionalGeneration.from_pretrained(MODEL_ID, trust_remote_code=True, torch_dtype=torch.float16).to("cuda").eval()
 
 
 
 
17
  processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
18
 
19
  DESCRIPTION = "[Qwen2-VL-7B Demo](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct)"
20
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
  @spaces.GPU
22
- def qwen_inference(media_path, text_input=None):
23
-
24
- image_extensions = Image.registered_extensions()
25
- if media_path.endswith(tuple([i for i, f in image_extensions.items()])):
26
- media_type = "image"
27
- elif media_path.endswith(("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")): # Check if it's a video path
28
- media_type = "video"
29
- else:
30
- raise ValueError("Unsupported media type. Please upload an image or video.")
31
-
 
 
 
 
 
 
 
 
 
 
32
  messages = [
33
  {
34
  "role": "user",
@@ -36,15 +83,17 @@ def qwen_inference(media_path, text_input=None):
36
  {
37
  "type": media_type,
38
  media_type: media_path,
39
- **({"fps": 8.0} if media_type == "video" else {}),
40
  },
41
  {"type": "text", "text": text_input},
42
  ],
43
  }
44
  ]
45
 
46
- text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
47
- image_inputs, video_inputs = process_vision_info(messages)
 
 
48
  inputs = processor(
49
  text=[text],
50
  images=image_inputs,
@@ -53,11 +102,18 @@ def qwen_inference(media_path, text_input=None):
53
  return_tensors="pt",
54
  ).to("cuda")
55
 
56
- generated_ids = model.generate(**inputs, max_new_tokens=1024)
57
- generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
58
- output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
59
-
60
- return output_text
 
 
 
 
 
 
 
61
 
62
  css = """
63
  #output {
@@ -73,12 +129,16 @@ with gr.Blocks(css=css) as demo:
73
  with gr.Tab(label="Image/Video Input"):
74
  with gr.Row():
75
  with gr.Column():
76
- input_media = gr.File(label="Upload Image or Video", type="filepath")
 
 
77
  text_input = gr.Textbox(label="Question")
78
  submit_btn = gr.Button(value="Submit")
79
  with gr.Column():
80
  output_text = gr.Textbox(label="Output Text")
81
 
82
- submit_btn.click(qwen_inference, [input_media, text_input], [output_text])
 
 
83
 
84
  demo.launch(debug=True)
 
1
  import gradio as gr
2
  import spaces
3
+ from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
4
  from qwen_vl_utils import process_vision_info
5
  import torch
6
  from PIL import Image
7
  import subprocess
8
  import numpy as np
9
  import os
10
+ from threading import Thread
11
+ import uuid
12
+ import io
13
 
14
  # Model and Processor Loading (Done once at startup)
15
  MODEL_ID = "Qwen/Qwen2-VL-7B-Instruct"
16
+ model = Qwen2VLForConditionalGeneration.from_pretrained(
17
+ MODEL_ID,
18
+ trust_remote_code=True,
19
+ torch_dtype=torch.float16
20
+ ).to("cuda").eval()
21
  processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
22
 
23
  DESCRIPTION = "[Qwen2-VL-7B Demo](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct)"
24
 
25
+ image_extensions = Image.registered_extensions()
26
+ video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")
27
+
28
+
29
+ def identify_and_save_blob(blob_path):
30
+ """Identifies if the blob is an image or video and saves it accordingly."""
31
+ try:
32
+ with open(blob_path, 'rb') as file:
33
+ blob_content = file.read()
34
+
35
+ # Try to identify if it's an image
36
+ try:
37
+ Image.open(io.BytesIO(blob_content)).verify() # Check if it's a valid image
38
+ extension = ".png" # Default to PNG for saving
39
+ media_type = "image"
40
+ except (IOError, SyntaxError):
41
+ # If it's not a valid image, assume it's a video
42
+ extension = ".mp4" # Default to MP4 for saving
43
+ media_type = "video"
44
+
45
+ # Create a unique filename
46
+ filename = f"temp_{uuid.uuid4()}_media{extension}"
47
+ with open(filename, "wb") as f:
48
+ f.write(blob_content)
49
+
50
+ return filename, media_type
51
+
52
+ except FileNotFoundError:
53
+ raise ValueError(f"The file {blob_path} was not found.")
54
+ except Exception as e:
55
+ raise ValueError(f"An error occurred while processing the file: {e}")
56
+
57
+
58
  @spaces.GPU
59
+ def qwen_inference(media_input, text_input=None):
60
+ if isinstance(media_input, str): # If it's a filepath
61
+ media_path = media_input
62
+ if media_path.endswith(tuple([i for i, f in image_extensions.items()])):
63
+ media_type = "image"
64
+ elif media_path.endswith(video_extensions):
65
+ media_type = "video"
66
+ else:
67
+ try:
68
+ media_path, media_type = identify_and_save_blob(media_input)
69
+ print(media_path, media_type)
70
+ except Exception as e:
71
+ print(e)
72
+ raise ValueError(
73
+ "Unsupported media type. Please upload an image or video."
74
+ )
75
+
76
+
77
+ print(media_path)
78
+
79
  messages = [
80
  {
81
  "role": "user",
 
83
  {
84
  "type": media_type,
85
  media_type: media_path,
86
+ **({"fps": 8.0} if media_type == "video" else {}),
87
  },
88
  {"type": "text", "text": text_input},
89
  ],
90
  }
91
  ]
92
 
93
+ text = processor.apply_chat_template(
94
+ messages, tokenize=False, add_generation_prompt=True
95
+ )
96
+ image_inputs, video_inputs = process_vision_info(messages)
97
  inputs = processor(
98
  text=[text],
99
  images=image_inputs,
 
102
  return_tensors="pt",
103
  ).to("cuda")
104
 
105
+ streamer = TextIteratorStreamer(
106
+ processor, skip_prompt=True, **{"skip_special_tokens": True}
107
+ )
108
+ generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
109
+
110
+ thread = Thread(target=model.generate, kwargs=generation_kwargs)
111
+ thread.start()
112
+
113
+ buffer = ""
114
+ for new_text in streamer:
115
+ buffer += new_text
116
+ yield buffer
117
 
118
  css = """
119
  #output {
 
129
  with gr.Tab(label="Image/Video Input"):
130
  with gr.Row():
131
  with gr.Column():
132
+ input_media = gr.File(
133
+ label="Upload Image or Video", type="filepath"
134
+ )
135
  text_input = gr.Textbox(label="Question")
136
  submit_btn = gr.Button(value="Submit")
137
  with gr.Column():
138
  output_text = gr.Textbox(label="Output Text")
139
 
140
+ submit_btn.click(
141
+ qwen_inference, [input_media, text_input], [output_text]
142
+ )
143
 
144
  demo.launch(debug=True)