update

README.md

@@ -1,6 +1,6 @@
 ---
 title: CapRL
-emoji: 🌖
+emoji: 🚀
 colorFrom: purple
 colorTo: green
 sdk: gradio
@@ -8,7 +8,7 @@ sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: Generate captions for images with CapRL (CPU-only)
+short_description: Generate captions for images with CapRL
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -7,6 +7,52 @@ MODEL_ID = "internlm/CapRL-3B"
 DEFAULT_PROMPT = "Describe the image in detail."
 MAX_NEW_TOKENS = 4096
 
+# Defaults for UI
+DEFAULT_IMAGE_PATH = "./examples/example_chinese.png"
+DEFAULT_CAPTION = """The image depicts a text excerpt, likely from an ancient or classical Chinese document, written in traditional script. Here is a detailed description:
+
+---
+
+**Image Description:**
+
+The image shows a page from what appears to be a historical or literary text, specifically related to ancient Chinese poetry or prose. The text is titled "序" (preface), indicating it's an introductory section. The content is structured as follows:
+
+1. **Header and Title:**
+   - The top line reads "古人云五百年萬六千日何況人生不閨雅佳" which translates to "Ancient people say that even five hundred years and six thousand days are nothing compared to life, not to mention being unreservedly beautiful."
+
+2. **Main Text Content:**
+   - The passage begins with "詎醒世之人萬六千日謂生人一事無成" suggesting a reflection on the fleeting nature of life over a vast period.
+
+3. **Key Phrases and Names:
+   - "過道還書身易張日夕思慕最凡在秋神冷" refers to someone who frequently thinks about books and autumn, possibly a scholar or writer.
+
+4. **Detailed Descriptions:
+   - "道士六魌急力問夢至春田處於筆頭翻得詩" describes a道士在六个地狱中匆忙询问梦境，最终在春天的田野处从笔下翻出诗句。 This indicates a dream or mystical experience leading to poetic inspiration.
+
+5. **Cultural References:
+   - "曲數本其間四時風景幽蘭處蜘蛛繚繞之思應如知" mentions counting the four seasons and the beauty of幽兰，with spiders weaving around, suggesting a vivid, natural scene.
+
+6. **Time and Weather:
+   - "在日前不眠眼中多釃精橋琥珀色忘冰奚與異" indicates someone who was awake all night, observing and noting the color of a bridge made of精 (possibly a type of stone or material) and ice, perhaps reflecting on a specific moment or event.
+
+7. **Emotional and Personal Touch:
+   - "悉心之樂不禁語話諸兄以初意手" expresses deep personal joy and initial intentions, likely referring to a heartfelt message or letter from someone named "初意手."
+
+8. **Additional Notes:
+   - "錄載曰亦自作不由得清退法這法後合因人省" suggests that the text is a record of someone's own writing, perhaps a self-reflection or diary entry, mentioning a need for self-retirement or withdrawal from certain matters.
+
+9. **Numerical References:
+   - "萬人云五百年" emphasizes the vastness of time, stating five hundred years.
+
+10. **Characters Mentioned:**
+    - The text mentions "鍾戴" (Zhòng dài), likely a person or place name, who is described as not leaving at a certain time, implying they were engaged in some activity until late.
+
+11. **Visual Elements:**
+    - The text is written in traditional Chinese characters, with a formal and classical style typical of ancient literary works.
+
+This detailed description should provide a pure text model with sufficient context to answer any related questions about the image."""
+DEFAULT_TOKENS = 674
+
 
 def load_model():
     device = "cpu"
@@ -14,7 +60,7 @@ def load_model():
 
     model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         MODEL_ID,
-        torch_dtype=dtype,
+        dtype=dtype,
         device_map="cpu",
         trust_remote_code=True,
         low_cpu_mem_usage=True,
@@ -28,7 +74,7 @@ MODEL, PROCESSOR = load_model()
 
 
 @torch.inference_mode()
-def generate_caption(image: Image.Image):
+def generate_caption(image: Image.Image, max_new_tokens: int = MAX_NEW_TOKENS):
     if image is None:
         return "", 0
 
@@ -61,9 +107,15 @@ def generate_caption(image: Image.Image):
             return_tensors="pt",
         ).to(device)
 
+        # Ensure slider value is an integer within bounds
+        try:
+            max_tokens = int(max(32, min(4096, int(max_new_tokens))))
+        except Exception:
+            max_tokens = MAX_NEW_TOKENS
+
         generated_ids = MODEL.generate(
             **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
+            max_new_tokens=max_tokens,
             do_sample=False,
         )
 
@@ -102,43 +154,56 @@ with gr.Blocks(title="CapRL Image Captioning (CPU)") as demo:
     )
 
     gr.Markdown(
-        "👉 Prefer faster inference? Try the GPU Space: "
-        "<a href=\"https://huggingface.co/spaces/yuhangzang/caprl\">yuhangzang/caprl</a>"
+        """
+<div style="font-size: 1.2rem; font-weight: 800; color: #e67300;">
+👉 Prefer faster inference? Try the GPU Space:
+<a href="https://huggingface.co/spaces/yuhangzang/caprl" style="color: #e67300; text-decoration: underline; font-weight: 900;">
+caprl (GPU Space)
+</a>
+</div>
+"""
     )
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="Input Image")
+            image_input = gr.Image(value=DEFAULT_IMAGE_PATH, type="pil", label="Input Image")
+            max_new_tokens_slider = gr.Slider(
+                minimum=32,
+                maximum=4096,
+                step=1,
+                value=MAX_NEW_TOKENS,
+                label="Max New Tokens (32–4096)",
+            )
             generate_button = gr.Button("Generate Caption")
         with gr.Column():
-            caption_output = gr.Textbox(label="Caption", lines=6)
-            token_output = gr.Number(label="Generated Tokens", precision=0)
+            caption_output = gr.Textbox(value=DEFAULT_CAPTION, label="Caption", lines=6)
+            token_output = gr.Number(value=DEFAULT_TOKENS, label="Generated Tokens", precision=0)
 
     generate_button.click(
         fn=generate_caption,
-        inputs=image_input,
+        inputs=[image_input, max_new_tokens_slider],
         outputs=[caption_output, token_output],
         show_progress=True,
     )
 
     image_input.upload(
         fn=generate_caption,
-        inputs=image_input,
+        inputs=[image_input, max_new_tokens_slider],
         outputs=[caption_output, token_output],
         show_progress=True,
     )
 
     gr.Examples(
         examples=[
-            ["./examples/example_chinese.png"],
-            ["./examples/example_receipt.jpg"],
-            ["./examples/example_table.png"],
+            ["./examples/example_chinese.png", MAX_NEW_TOKENS],
+            ["./examples/example_receipt.jpg", MAX_NEW_TOKENS],
+            ["./examples/example_table.png", MAX_NEW_TOKENS],
         ],
-        inputs=image_input,
+        inputs=[image_input, max_new_tokens_slider],
         outputs=[caption_output, token_output],
         fn=generate_caption,
         cache_examples=True,
-        label="📸 Example Images"
+        label="📸 Example Images",
     )
 
     gr.Markdown("### Citation")