File size: 11,436 Bytes
a4ba82a | 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 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | from collections import Counter
from pathlib import Path
import gc
import hashlib
import json
import os
import threading
import time
from datetime import datetime, timezone
from .config import MODELS, VENDOR, MODEL_CACHE
_LOCK = threading.Lock()
_MANIFEST_LOCK = threading.Lock()
_VERIFIED_MANIFESTS = {}
class NotReady(RuntimeError):
pass
def verify_cuda_residency(model):
devices = Counter((str(p.device) for p in model.parameters()))
if not devices or any((not d.startswith('cuda:') for d in devices)):
raise NotReady('GPU memory is insufficient for this native Captioner configuration.')
return dict(devices)
def runtime_event(event, **values):
print('FLORA_RUNTIME ' + json.dumps(dict(event=event, utc=datetime.now(timezone.utc).isoformat(), **values), sort_keys=True), flush=True)
def ready():
if os.environ.get('FLORA_ENABLE_INFERENCE', '0') != '1':
return (False, 'Preview ready 路 live annotation is paused.')
import torch
if not torch.cuda.is_available():
return (False, 'Preview ready 路 GPU hardware is needed for annotation.')
if sum((torch.cuda.get_device_properties(i).total_memory for i in range(torch.cuda.device_count()))) < 75 * 1024 ** 3:
return (False, 'Annotation needs an 80 GB GPU or a validated multi-GPU configuration.')
if not os.environ.get('HF_TOKEN'):
return (False, 'Annotation is waiting for access to the private weights.')
return (True, 'Ready to annotate 路 one video at a time.')
def snapshot(kind):
from huggingface_hub import snapshot_download
spec = dict(MODELS[kind])
spec['repo'] = os.environ.get(spec.get('repo_env', ''), spec['repo'])
if not spec['repo']:
raise ValueError('Set ' + spec['repo_env'] + ' to the released checkpoint repository.')
patterns = ['adapter_config.json', 'adapter_model.safetensors', 'config.json', 'generation_config.json', 'model.safetensors.index.json', 'model-*.safetensors', 'processor_config.json', 'tokenizer.json', 'tokenizer_config.json', 'chat_template.jinja', 'README.md', 'MODEL_PROVENANCE.json', 'SOURCE_MANIFEST.sha256']
return Path(snapshot_download(spec['repo'], revision=spec['revision'], token=os.environ.get('HF_TOKEN'), allow_patterns=patterns, local_dir=MODEL_CACHE / kind))
def processor(kind, path):
from transformers import AutoProcessor
from transformers.models.qwen3_vl.video_processing_qwen3_vl import Qwen3VLVideoProcessor
proc = AutoProcessor.from_pretrained(path, local_files_only=True, use_fast=True, min_pixels=1024, max_pixels=589824, video_processor=Qwen3VLVideoProcessor.from_pretrained(path, local_files_only=True, fps=4 if kind == 'tempseg' else 8, max_frames=128 if kind == 'tempseg' else 40))
proc.chat_template = (path / 'chat_template.jinja').read_text()
return proc
def validate_manifest(path):
with _MANIFEST_LOCK:
_validate_manifest(path)
def _validate_manifest(path):
started = time.monotonic()
allowed_root = path.parent.parent if path.parent.name == 'snapshots' else path
manifest = (path / 'SOURCE_MANIFEST.sha256').read_bytes()
key = str(path.resolve())
manifest_sha = hashlib.sha256(manifest).hexdigest()
files = []
def fingerprint(source):
resolved = source.resolve()
if not resolved.is_relative_to(allowed_root.resolve()):
raise RuntimeError('Invalid model manifest')
stat = resolved.stat()
return (str(resolved), stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns)
for line in manifest.decode().splitlines():
if not line.strip():
continue
expected, filename = line.split(maxsplit=1)
relative = Path(filename.lstrip('*'))
if relative.is_absolute() or '..' in relative.parts:
raise RuntimeError('Invalid model manifest')
source = path / relative
if not source.resolve().is_relative_to(allowed_root.resolve()):
runtime_event('manifest_path_rejected', file=filename, snapshot=str(path), target=str(source.resolve()), allowed_root=str(allowed_root.resolve()))
raise RuntimeError('Invalid model manifest')
files.append((source, expected, fingerprint(source)))
before = tuple(((str(source), expected, stamp) for source, expected, stamp in files))
if _VERIFIED_MANIFESTS.get(key) == (manifest_sha, before):
runtime_event('manifest_verified_reuse', kind=path.name, files=len(files), payload_bytes_read=0, elapsed_s=round(time.monotonic() - started, 3))
return
_VERIFIED_MANIFESTS.pop(key, None)
for source, expected, _ in files:
with source.open('rb') as f:
if hashlib.file_digest(f, 'sha256').hexdigest() != expected:
raise RuntimeError('A model asset failed its integrity check')
after = tuple(((str(source), expected, fingerprint(source)) for source, expected, _ in files))
if before != after or manifest != (path / 'SOURCE_MANIFEST.sha256').read_bytes():
raise RuntimeError('A model asset changed during its integrity check')
_VERIFIED_MANIFESTS[key] = (manifest_sha, after)
runtime_event('manifest_full_verification', kind=path.name, files=len(files), payload_bytes_read=sum((stamp[3] for _, _, stamp in files)), elapsed_s=round(time.monotonic() - started, 3))
def preflight_adapter_download():
path = snapshot('tempseg')
validate_manifest(path)
runtime_event('adapter_download_preflight_pass', revision=MODELS['tempseg']['revision'], adapter_sha256=MODELS['tempseg']['sha256'], directory=str(path))
class Session:
def __init__(self):
self.model = self.proc = None
self.audit = {}
def load(self, kind):
import torch
previous_threads = torch.get_num_threads()
torch.set_num_threads(4)
load_started = time.monotonic()
from transformers import AutoModelForImageTextToText
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel
from production_gdn_fix import _keep_native_gdn_fp32
from verify_gdn import verify
self.close()
runtime_event('model_load_start', kind=kind, cpu_threads=torch.get_num_threads(), previous_cpu_threads=previous_threads)
_keep_native_gdn_fp32(Qwen3_5PreTrainedModel)
path = snapshot(kind)
runtime_event('model_assets_ready', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
validate_manifest(path)
runtime_event('model_assets_verified', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
if kind == 'tempseg':
base_spec = MODELS['tempseg_base']
from huggingface_hub import snapshot_download
base = snapshot_download(base_spec['repo'], revision=base_spec['revision'], token=os.environ.get('HF_TOKEN'), allow_patterns=['*.safetensors', '*.json', '*.jinja'])
from tempseg_fp32_loader import load
records = {}
self.model = load({'base': base, 'export': str(path), 'adapter_sha256': MODELS[kind]['sha256']}, lambda name, value: records.__setitem__(name, value))
self.audit[kind] = records
else:
index = json.loads((path / 'model.safetensors.index.json').read_text())
extra = {k for k in index['weight_map'] if '_extra_state' in k}
memory = {i: f'{max(1, int(torch.cuda.get_device_properties(i).total_memory / 1024 ** 3) - 12)}GiB' for i in range(torch.cuda.device_count())}
self.model, info = AutoModelForImageTextToText.from_pretrained(path, local_files_only=True, dtype=torch.bfloat16, attn_implementation={'text_config': 'sdpa', 'vision_config': 'sdpa'}, device_map='balanced', max_memory=memory, output_loading_info=True)
if any((info.get(k) for k in ['missing_keys', 'mismatched_keys', 'error_msgs'])) or not set(info.get('unexpected_keys', [])) <= extra:
raise RuntimeError('The released Captioner did not load completely')
devices = verify_cuda_residency(self.model)
self.model.eval()
self.audit[kind] = {'gdn': verify(self.model, path), 'loading_info': info, 'parameter_devices': devices, 'parameter_dtypes': dict(Counter((str(p.dtype) for p in self.model.parameters())))}
runtime_event('model_parameters_verified', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
self.proc = processor(kind, path)
if kind == 'tempseg':
adapter = self.audit[kind]['ADAPTER_VALUE_VERIFICATION.json']
gdn = self.audit[kind]['GDN_VALUE_VERIFICATION.json']
verification = dict(adapter_tensors=adapter['tensors'], adapter_values_exact=adapter['all_saved_values_exact_before_and_after_placement'], adapter_fp32=adapter['holders_fp32_before_assignment'], gdn_values_exact=gdn['all_values_equal_saved_checkpoint'], gdn_fp32=gdn['all_runtime_fp32'])
else:
gdn = self.audit[kind]['gdn']
verification = dict(gdn_values_exact=gdn['all_values_equal_saved_checkpoint'], gdn_fp32=gdn['all_runtime_fp32'], parameter_dtypes=self.audit[kind]['parameter_dtypes'])
runtime_event('model_ready', kind=kind, revision=MODELS[kind]['revision'], torch_version=torch.__version__, gpu=torch.cuda.get_device_name(0), allocated_gib=round(torch.cuda.memory_allocated() / 1024 ** 3, 3), **verification)
return (self.model, self.proc)
def close(self):
self.model = self.proc = None
gc.collect()
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
def generate(model, proc, inputs, budget):
import torch
inputs = dict(inputs)
inputs.pop('video_metadata', None)
n = inputs['input_ids'].shape[1]
if n + budget > 32768:
raise InputErrorForModel('The selected video exceeds the native model context budget.')
input_tensors = {k: {'dtype': str(v.dtype), 'shape': list(v.shape), 'sha256': hashlib.sha256(v.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()} for k, v in inputs.items() if torch.is_tensor(v)}
input_tensor_digest = hashlib.sha256(json.dumps(input_tensors, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
device = next(model.parameters()).device
args = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
eos = proc.tokenizer.eos_token_id
started = time.monotonic()
runtime_event('generation_start', input_tokens=n, output_limit=budget, input_tensor_digest=input_tensor_digest, input_tensors=input_tensors)
with torch.inference_mode():
output = model.generate(**args, do_sample=False, num_beams=1, max_new_tokens=budget, eos_token_id=eos, pad_token_id=proc.tokenizer.pad_token_id, use_cache=True, max_time=300)
ids = output[0, n:].cpu().tolist()
stop = bool(ids) and ids[-1] in ([eos] if isinstance(eos, int) else eos)
text = proc.tokenizer.decode(ids, skip_special_tokens=True).strip()
runtime_event('generation_end', input_tokens=n, output_limit=budget, generated_tokens=len(ids), complete=stop and bool(text), elapsed_s=round(time.monotonic() - started, 3), peak_allocated_gib=round(torch.cuda.max_memory_allocated() / 1024 ** 3, 3))
if not stop or not text:
raise InputErrorForModel('The model did not return a complete output for this input.')
return text
class InputErrorForModel(RuntimeError):
pass
|