Upload app.py with huggingface_hub

app.py

@@ -5,7 +5,7 @@ from queue import Queue
 import threading
 import torch
 from PIL import Image, ImageOps
-from transformers import AutoProcessor, MiniCPMV4_6ForConditionalGeneration, TextIteratorStreamer
+from transformers import AutoProcessor, MiniCPMV4_6ForConditionalGeneration
 import gradio as gr
 
 try:
@@ -88,6 +88,9 @@ def env_flag(name: str, default: bool = False) -> bool:
     return value.strip().lower() in {"1", "true", "yes", "on"}
 
 
+ENABLE_MODEL_WARMUP = env_flag("NOTEWORTHY_WARMUP", True)
+
+
 def _order_points(points):
     import numpy as np
 
@@ -294,54 +297,19 @@ print("Models loaded.")
 def stream_model(model, image: Image.Image):
     messages = [{
         "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": NOTES_PROMPT},
-        ],
+        "content": [image, NOTES_PROMPT],
     }]
 
     with torch.no_grad():
-        inputs = processor.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=True,
-            return_dict=True,
-            return_tensors="pt",
-            enable_thinking=False,
-            processor_kwargs={
-                "downsample_mode": "4x",
-                "max_slice_nums": 9,
-                "use_image_id": True,
-            },
-        ).to(model.device)
-
-    for k, v in inputs.items():
-        if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
-            inputs[k] = v.to(dtype=torch.bfloat16)
-
-    streamer = TextIteratorStreamer(
-        processor.tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True,
-    )
-
-    thread = threading.Thread(
-        target=model.generate,
-        kwargs={
-            **inputs,
-            "max_new_tokens": 1024,
-            "do_sample": False,
-            "num_beams": 1,
-            "streamer": streamer,
-            "downsample_mode": "4x",
-        },
-    )
-    thread.start()
-
-    for chunk in streamer:
-        yield chunk
+        result = model.chat(
+            image=None,
+            msgs=messages,
+            tokenizer=processor.tokenizer,
+            sampling=False,
+            max_new_tokens=1024,
+        )
 
-    thread.join()
+    yield result or ""
 
 
 def stream_model_text(model, image: Image.Image):
@@ -351,6 +319,47 @@ def stream_model_text(model, image: Image.Image):
         yield text
 
 
+def warmup_models():
+    if not ENABLE_MODEL_WARMUP:
+        print("Model warmup disabled.")
+        return
+
+    warmup_path = "examples/000100005-1_1_1.png"
+    if not os.path.exists(warmup_path):
+        print("Skipping model warmup; example image is missing.")
+        return
+
+    print("Warming up local models...")
+    image, processed_image_path = preprocess_sheet_music_image(warmup_path)
+    try:
+        messages = [{
+            "role": "user",
+            "content": [image, NOTES_PROMPT],
+        }]
+        for name, model in (
+            ("fine-tuned", finetuned_model),
+            ("original", original_model),
+        ):
+            print(f"  Warming {name} model...")
+            with torch.no_grad():
+                model.chat(
+                    image=None,
+                    msgs=messages,
+                    tokenizer=processor.tokenizer,
+                    sampling=False,
+                    max_new_tokens=8,
+                )
+        print("Model warmup complete.")
+    finally:
+        try:
+            os.unlink(processed_image_path)
+        except OSError:
+            pass
+
+
+warmup_models()
+
+
 def _stream_openai_transcription(client, model_id, mime, image_data):
     response_stream = client.chat.completions.create(
         model=model_id,
@@ -547,7 +556,6 @@ with gr.Blocks(
         type="pil",
         label="Preprocessed Image Sent to All Models",
         interactive=False,
-        show_download_button=True,
     )
     gr.Examples(
         examples=[