prithivMLmods commited on
Commit
dab83dd
·
verified ·
1 Parent(s): 7799820

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -43,6 +43,8 @@ model_m = AutoModel.from_pretrained(
43
  trust_remote_code=True,
44
  torch_dtype=torch.float16
45
  ).to(device).eval()
 
 
46
 
47
  # Load Space Thinker
48
  MODEL_ID_Z = "remyxai/SpaceThinker-Qwen2.5VL-3B"
@@ -93,7 +95,7 @@ def generate_image(model_name: str, text: str, image: Image.Image,
93
  if image is None:
94
  yield "Please upload an image."
95
  return
96
- # Construct message with <image> token as per reference
97
  if "<image>" not in text:
98
  message = f"<image>\n{text}"
99
  else:
@@ -346,7 +348,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
346
  model_choice = gr.Radio(
347
  choices=["Llama-3.1-Nemotron-Nano-VL-8B-V1", "SpaceThinker-3B", "coreOCR-7B-050325-preview"],
348
  label="Select Model",
349
- value="Llama-3.1-Nemotron-Nano-VL-8B-V1" # Updated default value to a valid choice
350
  )
351
 
352
  gr.Markdown("**Model Info**")
 
43
  trust_remote_code=True,
44
  torch_dtype=torch.float16
45
  ).to(device).eval()
46
+ # Fix AssertionError by setting img_context_token_id
47
+ model_m.img_context_token_id = tokenizer_m.convert_tokens_to_ids("<image>")
48
 
49
  # Load Space Thinker
50
  MODEL_ID_Z = "remyxai/SpaceThinker-Qwen2.5VL-3B"
 
95
  if image is None:
96
  yield "Please upload an image."
97
  return
98
+ # Construct message with <image> token
99
  if "<image>" not in text:
100
  message = f"<image>\n{text}"
101
  else:
 
348
  model_choice = gr.Radio(
349
  choices=["Llama-3.1-Nemotron-Nano-VL-8B-V1", "SpaceThinker-3B", "coreOCR-7B-050325-preview"],
350
  label="Select Model",
351
+ value="Llama-3.1-Nemotron-Nano-VL-8B-V1"
352
  )
353
 
354
  gr.Markdown("**Model Info**")