Spaces:
Runtime error
Runtime error
File size: 5,282 Bytes
4a02afe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Modal app: runs MiniCPM-V-4.6 captioning on a remote GPU."""
import modal
app = modal.App("photographers-archive")
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"transformers[torch]>=5.7.0",
"torchvision",
"av",
"Pillow",
"torch>=2.1.0",
"accelerate",
)
)
MODEL_ID = "openbmb/MiniCPM-V-4.6"
CAPTION_PROMPT = """You are a wedding and portrait photo archivist. Analyze this image and return ONLY a valid JSON object — no markdown, no explanation, no code fences.
Crucial Prompting Guidelines:
1. Be highly specific with textures, fabrics, and patterns in "attire" (e.g., "lace wedding gown", "black velvet tuxedo", "pinstripe suit").
2. Avoid generic descriptions in "summary". Focus on explicit visual facts (who, what, where, explicit actions).
3. Under "primary_subjects", explicitly label roles if evident (e.g., "bride", "groom", "bridesmaid", "groomsman", "mother of the bride").
4. For "depth_of_field", specify features like "shallow depth of field", "bokeh background", or "deep focus".
Use this exact schema:
{
"summary": "2-3 sentence description of the scene",
"subjects": {
"people_count": 0,
"primary_subjects": [],
"relationships": [],
"attire": []
},
"scene": {
"location_type": "",
"environment": "",
"setting_details": []
},
"actions": {
"primary_action": "",
"body_language": []
},
"lighting": {
"lighting_style": "",
"time_of_day_estimate": ""
},
"composition": {
"shot_type": "",
"camera_angle": ""
},
"mood": {
"primary_emotions": [],
"atmosphere": ""
},
"technical_cues": {
"color_palette": [],
"depth_of_field": ""
},
"search_tags": [],
"archive_keywords": []
}"""
# Cache model weights in a Modal Volume so they persist across cold starts
model_volume = modal.Volume.from_name("minicpm-weights", create_if_missing=True)
@app.cls(
image=image,
gpu="A10G",
volumes={"/model-cache": model_volume},
timeout=300,
)
class Captioner:
@modal.enter()
def load(self):
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
self.processor = AutoProcessor.from_pretrained(
MODEL_ID, cache_dir="/model-cache"
)
self.model = AutoModelForImageTextToText.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
cache_dir="/model-cache",
)
@modal.method()
def caption(self, image_bytes: bytes, filename: str = "image.jpg") -> str:
import json
import os
import re
import tempfile
import torch
suffix = os.path.splitext(filename)[-1] or ".jpg"
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
f.write(image_bytes)
tmp_path = f.name
try:
messages = [
{
"role": "user",
"content": [
{"type": "image", "url": tmp_path},
{"type": "text", "text": CAPTION_PROMPT},
],
}
]
inputs = self.processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
downsample_mode="16x",
max_slice_nums=9,
).to(self.model.device)
with torch.inference_mode():
generated_ids = self.model.generate(
**inputs,
downsample_mode="16x",
max_new_tokens=2048,
do_sample=False,
)
trimmed = [
out_ids[len(in_ids):]
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
raw = self.processor.batch_decode(
trimmed,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0].strip()
# Strip markdown code fences if model added them anyway
raw = re.sub(r"^```(?:json)?\s*", "", raw)
raw = re.sub(r"\s*```$", "", raw)
# Validate JSON — if broken, return as plain text so it still gets stored
try:
json.loads(raw)
except json.JSONDecodeError:
# Attempt to salvage by extracting the outermost {...} block
match = re.search(r"\{.*\}", raw, re.DOTALL)
if match:
candidate = match.group(0)
try:
json.loads(candidate)
return candidate
except json.JSONDecodeError:
pass
# Give up and return raw — caption_store will treat it as plain text
return raw
return raw
finally:
os.unlink(tmp_path)
|