JimmyK300 commited on
Commit
534e4f4
·
verified ·
1 Parent(s): 1124359

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -6
app.py CHANGED
@@ -7,11 +7,13 @@ from pathlib import Path
7
  from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, Qwen2VLForConditionalGeneration
8
  from PIL import Image
9
 
 
 
10
  # Load Vision-Language Model
11
  vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
12
  "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto"
13
  )
14
- vl_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", max_pixels=240*240)
15
 
16
  # Load Text Model
17
  model_name = "Qwen/Qwen2.5-Math-1.5B-Instruct"
@@ -22,12 +24,15 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
22
  math_messages = []
23
 
24
  def resize_image(image):
25
- max_size = 480
26
  if isinstance(image, str): # Handle file paths
27
  image = Image.open(image)
28
- image.thumbnail((max_size, max_size))
29
- return image
30
-
 
 
 
 
31
  def process_image(image, shouldConvert=False):
32
  global math_messages
33
  math_messages = [] # Reset when uploading an image
@@ -96,7 +101,7 @@ with gr.Blocks(css=css) as demo:
96
  with gr.Column():
97
  with gr.Tabs() as input_tabs:
98
  with gr.Tab("Upload"):
99
- input_image = resize_image(gr.Image(type="pil", label="Upload"))
100
  with gr.Tab("Sketch"):
101
  input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
102
  input_tabs.select(fn=tabs_select, inputs=[state])
 
7
  from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, Qwen2VLForConditionalGeneration
8
  from PIL import Image
9
 
10
+ max_size = 240
11
+
12
  # Load Vision-Language Model
13
  vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
14
  "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto"
15
  )
16
+ vl_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", max_pixels=max_size*max_size)
17
 
18
  # Load Text Model
19
  model_name = "Qwen/Qwen2.5-Math-1.5B-Instruct"
 
24
  math_messages = []
25
 
26
  def resize_image(image):
 
27
  if isinstance(image, str): # Handle file paths
28
  image = Image.open(image)
29
+ try:
30
+ image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
31
+ return image
32
+ except Exception as e:
33
+ print(f"Error resizing image: {e}")
34
+ return None
35
+
36
  def process_image(image, shouldConvert=False):
37
  global math_messages
38
  math_messages = [] # Reset when uploading an image
 
101
  with gr.Column():
102
  with gr.Tabs() as input_tabs:
103
  with gr.Tab("Upload"):
104
+ input_image = gr.Image(type="pil", label="Upload")
105
  with gr.Tab("Sketch"):
106
  input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
107
  input_tabs.select(fn=tabs_select, inputs=[state])