Compare commits

...
10 changed files with 658 additions and 0 deletions
+9
View File
@@ -126,7 +126,15 @@ except Exception as _e: # noqa
print(f"[BFSNodes] MiniMax H3 CLIP video context not loaded: {_e!r}")
H3CLIP_NODE_CLASS_MAPPINGS, H3CLIP_NODE_DISPLAY_NAME_MAPPINGS = {}, {}
try:
from .minimax_h3_embedding_bridge import NODE_CLASS_MAPPINGS as H3BRIDGE_NODE_CLASS_MAPPINGS
from .minimax_h3_embedding_bridge import NODE_DISPLAY_NAME_MAPPINGS as H3BRIDGE_NODE_DISPLAY_NAME_MAPPINGS
except Exception as _e: # noqa
print(f"[BFSNodes] Experimental H3 embedding bridge not loaded: {_e!r}")
H3BRIDGE_NODE_CLASS_MAPPINGS, H3BRIDGE_NODE_DISPLAY_NAME_MAPPINGS = {}, {}
NODE_CLASS_MAPPINGS = {
**H3BRIDGE_NODE_CLASS_MAPPINGS,
**H3CLIP_NODE_CLASS_MAPPINGS,
**H3GUIDE_NODE_CLASS_MAPPINGS,
**SHOT_NODE_CLASS_MAPPINGS,
@@ -157,6 +165,7 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
**H3BRIDGE_NODE_DISPLAY_NAME_MAPPINGS,
**H3CLIP_NODE_DISPLAY_NAME_MAPPINGS,
**H3GUIDE_NODE_DISPLAY_NAME_MAPPINGS,
**SHOT_NODE_DISPLAY_NAME_MAPPINGS,
+133
View File
@@ -0,0 +1,133 @@
# Experimental Qwen3.5 9B thinking states -> H3 conditioning
Branch: `experiment/h3-9b-thinking-embeddings`. No changes to main or the live
OpenJutsu training run. This is a prepared research prototype, not a trained
encoder replacement. GPU execution and real workflow validation are pending.
## What is actually tested
The full Qwen3.5-9B sees numbered pictures and chronological guide frames, plus
an unchanged copy of the request. It generates with thinking off/on. A forward
hook captures its last decoder-block activations BEFORE final normalization:
one prefill bank and a separate bank of generated-token states. The sampled
last token may not have been forwarded. Reasoning text is not rewritten into an
H3 prompt, not saved, and not fed through the teacher. Token states themselves
are the bridge input. A completed `</think>` is required in thinking mode;
truncated thinking is not silently labelled successful.
A shared trainable bridge consumes these states. Position/modality queries
produce a variable-length sequence with the original H3 teacher's exact layout,
5120-wide unnormalized features and visual/text tags. The bridge never sees
teacher feature values as INPUT. Only its reconstruction loss uses them.
The 9B and H3 encoders stay frozen; H3 DiT/VAE are not loaded for calibration.
Teacher is the existing converted H3 encoder (50 layers), so no 66-GB full
teacher download is needed.
This prototype requires a cached teacher layout for each evaluation case. It is
NOT yet a general online encoder, a speedup demonstration, or an identity/motion
quality improvement. Distilling ordinary H3 embeddings may suppress useful
thinking information. Embedding error alone cannot prove thinking benefits.
## What is ready and what is not
Prepared: disjoint input manifest, teacher/student cache scripts, a modality-
balanced shared resampling bridge, train/validation/test reconstruction metrics,
held-out conditioning exports, and a guarded Comfy cached-case embedding swap.
Only syntax and CPU file preflight are checked. No tensor, GPU, generation or
quality tests have been run, per the user's local testing preference. No bridge
weights have been trained/exported. No new model download or environment upgrade
is performed automatically. GPU entrypoints refuse to run while the dataset
reviewer is active or GPU memory is occupied; there is no auto-start watcher.
Full local Qwen3.5-9B weights and a compatible Transformers environment are
required. Existing Comfy caption-model support does not guarantee a matching
Transformers model directory. Keep the experimental environment separate from
Comfy/training; use current official Qwen requirements rather than upgrading
the active environments. Generation uses a complete HF model, not GGUF and not
the headless H3 encoder.
## CPU preparation (safe during dataset processing)
```bash
cd /data/bfs-dev-vlm-reasoning/experiments/h3_embedding_bridge
python make_manifest.py --output /data/reports/h3_9b_embedding_bridge/manifest.jsonl
python preflight.py --manifest /data/reports/h3_9b_embedding_bridge/manifest.jsonl \
--output /data/reports/h3_9b_embedding_bridge/preflight.json
```
Default: 32 calibration-train, 8 validation and 8 held-out cases; no exact
reference-asset hash reuse across these cases. Inputs are original training
split captions/guides, capped at four pictures for the initial cost experiment.
They are UNFILTERED calibration requests, NOT approved H3 training examples and
NOT a clean identity-quality benchmark. No targets, original captions or
production caches are modified. Their source captions may be inaccurate; this
limits downstream identity conclusions. Reference-image hashes prevent exact
asset overlap, not all appearances of the same real identity.
Each probe uses at most 73 native 24-fps frames (~3.04 s), source-time resampling,
384-pixel width and 2-fps Qwen context. Teacher and student use the same media
interval and picture order, with their respective native vision processors.
This small probe cannot establish full-video identity stability.
## Deferred GPU stages (after the dataset reviewer exits)
Use absolute paths or set `PROBE`, `CALIB`, `H3PY`, `QWENPY` in your shell.
`H3PY` is a compatible Comfy Python environment; `QWENPY` is an isolated
Transformers environment supporting the complete Qwen3.5 model. Use an empty
output directory. Models are loaded sequentially, not concurrently.
```bash
PROBE=/data/bfs-dev-vlm-reasoning/experiments/h3_embedding_bridge
CALIB=/data/reports/h3_9b_embedding_bridge
H3PY=/data/comfyui/.venv/bin/python
# Set QWENPY to your separate compatible environment, and QWEN_MODEL to local weights.
PYTHONPATH=/tmp/caption-runtime-site "$H3PY" "$PROBE/cache_teacher.py" \
--manifest "$CALIB/manifest.jsonl" --cache "$CALIB/cache"
"$QWENPY" "$PROBE/cache_student.py" --model "$QWEN_MODEL" \
--manifest "$CALIB/manifest.jsonl" --cache "$CALIB/cache" --max-new-tokens 512
"$QWENPY" "$PROBE/train_bridge.py" --manifest "$CALIB/manifest.jsonl" \
--cache "$CALIB/cache" --output "$CALIB/run_001" --steps 500
```
Cache only complete, matching input keys. Teacher weights stay in their native
conversion/quantization format. If the local Comfy Python cannot load that format,
fix the isolated runtime first. The teacher stage has not been executed here.
Student HF paths are local-only; partial thinking halts with a metadata record
and can be retried with a larger token budget. The scripts do not resume training
optimizer state; choose a new output directory for each trial.
Validation selects one shared bridge using both thinking modes. Test outputs
contain `thinking_0`, `thinking_1`, and `thinking_ablated` (same thinking-on
prefill, generated states zeroed, same state-bank quota). Never use test teacher
features to fit the bridge. The per-channel output scale is estimated from the
training split only. There is no video/identity/reward objective in this first
calibration. Production use might require H3 adaptation, broader calibration,
and a trained task-specific objective.
## Local Comfy experiment after calibration
Copy held-out exported `*.thinking_*.safetensors` and
`*.thinking_ablated.safetensors` into `ComfyUI/models/h3_embedding_bridge/`.
Do not select `bridge.safetensors` there: it contains weights, not conditioning.
Restart Comfy with this experimental BFS branch. The node is:
**BFS H3 Experimental 9B Embedding Bridge (Cached Case)**
First build native H3 conditioning with the exact SAME encoder, prompt, ordered
pictures, source interval/resampling, dimensions and 2-fps guide context used in
that export. Feed its `positive` through the experimental node before the
sampler; retain the normal latent/Add Guide/reference-VAE path. The node replaces
only embeddings and retains existing metadata, native reference latents and
aligned guide latents. A teacher-baseline comparison rejects mismatched cached
cases/tags; do not raise tolerance to bypass wrong inputs. This baseline pass
means the prototype does not save encoder compute.
Render four matched arms with identical seed, guide, references, LoRA, CFG,
steps and sampler: native teacher, 9B thinking off, 9B thinking on, generated-state
ablation. Judge wrong-person assignments, identity exchanges after crossings or
occlusions, reference wardrobe fidelity and motion separately. A difference in
output, or a lower embedding reconstruction error, does not establish a benefit.
Use curated longer multi-person clips before changing the production encoder.
Official Qwen model/config: https://huggingface.co/Qwen/Qwen3.5-9B
Original H3 encoder: https://huggingface.co/MiniMaxAI/MiniMax-H3/tree/main/FL2VA/text_encoder
+78
View File
@@ -0,0 +1,78 @@
"""Experimental 9B hidden-state -> H3 layer-50 bridge, not an encoder replacement yet."""
import math
from dataclasses import asdict, dataclass
import torch
from torch import nn
from torch.nn import functional as F
@dataclass
class BridgeConfig:
student_dim: int = 4096
teacher_dim: int = 5120
width: int = 512
heads: int = 8
source_slots: int = 512
query_chunk: int = 256
class HiddenStateBridge(nn.Module):
"""Queries use only target positions/modality tags, NEVER teacher features.
Teacher layout is needed for this cached-case prototype. Arbitrary-input
layout prediction is deliberately not claimed or implemented.
"""
def __init__(self, config=BridgeConfig()):
super().__init__()
self.config = config
self.source = nn.Sequential(nn.LayerNorm(config.student_dim), nn.Linear(config.student_dim, config.width))
self.modality = nn.Embedding(2, config.width)
self.attention = nn.MultiheadAttention(config.width, config.heads, batch_first=True)
self.norm = nn.LayerNorm(config.width)
self.output = nn.Linear(config.width, config.teacher_dim)
self.register_buffer('teacher_channel_scale', torch.ones(config.teacher_dim))
def position(self, length, device):
# Normalized positions do not depend on teacher feature values.
pos = torch.linspace(0, 1, length, device=device)[:, None]
freq = torch.exp(torch.arange(0, self.config.width, 2, device=device) * (-math.log(10000) / self.config.width))
angles = pos * 1000 * freq
return torch.stack((angles.sin(), angles.cos()), -1).flatten(-2)
def memory(self, prefix, generated):
# Separate quotas keep thinking states from being washed out by image tokens.
banks = [self.source(prefix)]
if generated.shape[0]:
banks.append(self.source(generated))
quota = max(1, self.config.source_slots // len(banks))
pooled = [F.adaptive_avg_pool1d(x.T[None], min(quota, len(x))).transpose(1, 2) for x in banks]
result = torch.cat(pooled, 1)
return result + self.position(result.shape[1], result.device)[None].to(result.dtype)
def forward(self, prefix, generated, target_tags):
assert prefix.ndim == generated.ndim == 2
assert prefix.shape[-1] == generated.shape[-1] == self.config.student_dim
assert target_tags.ndim == 1 and len(target_tags) > 0
memory = self.memory(prefix, generated)
queries = self.modality(target_tags.long()) + self.position(len(target_tags), target_tags.device).to(memory.dtype)
output = []
for start in range(0, len(queries), self.config.query_chunk):
q = queries[None, start:start + self.config.query_chunk]
x, _ = self.attention(q, memory, memory, need_weights=False)
output.append(self.output(self.norm(x + q))[0])
return torch.cat(output) * self.teacher_channel_scale
def config_dict(self):
return asdict(self.config)
def alignment_loss(prediction, teacher, tags):
# Equal modality weighting prevents the longer visual prefix dominating text.
scale = teacher.float().square().mean(0).sqrt().clamp_min(1e-3)
losses = []
for kind in (0, 1):
mask = tags == kind
if mask.any():
p, t = prediction[mask].float(), teacher[mask].float()
losses.append(((p - t) / scale).square().mean() + .1 * (1 - F.cosine_similarity(p, t).mean()))
return torch.stack(losses).mean()
@@ -0,0 +1,86 @@
"""GPU stage 2. Capture actual 9B decoder activations, not rewritten prompts.
A forward hook retains one selected layer only, avoiding generate's all-layer
hidden-state cache. No generated reasoning text is persisted or fed to H3.
"""
import argparse, json, sys
from pathlib import Path
from common import load_manifest, require_idle_gpu, visual_inputs
def decoder_layer(model):
for path in ('model.language_model.layers', 'language_model.layers', 'model.layers'):
obj = model
try:
for name in path.split('.'):
obj = getattr(obj, name)
return obj[-1]
except AttributeError:
pass
raise RuntimeError('Unsupported Transformers decoder layout; no hidden states captured')
def main():
ap = argparse.ArgumentParser(); ap.add_argument('--manifest', required=True); ap.add_argument('--cache', required=True)
ap.add_argument('--model', required=True, help='Local full Qwen3.5-9B Transformers directory; no automatic download')
ap.add_argument('--max-new-tokens', type=int, default=512)
args = ap.parse_args(); require_idle_gpu()
if not Path(args.model).is_dir():
raise RuntimeError('Full 9B weights are not present locally. Download in an isolated environment first.')
import numpy as np
import torch
from safetensors.torch import save_file
from transformers import AutoModelForImageTextToText, AutoProcessor
processor = AutoProcessor.from_pretrained(args.model, local_files_only=True)
model = AutoModelForImageTextToText.from_pretrained(args.model, dtype=torch.bfloat16, device_map='auto', local_files_only=True)
model.eval().requires_grad_(False)
layer = decoder_layer(model); cache = Path(args.cache); cache.mkdir(parents=True, exist_ok=True)
for row in load_manifest(args.manifest):
refs, frames, picks, info = visual_inputs(row)
content = []
for i, image in enumerate(refs, 1):
content.extend([{'type': 'text', 'text': f'<Picture {i}>:'}, {'type': 'image', 'image': image}])
content.append({'type': 'text', 'text': '<Video 1> chronological motion guide:'})
for i in picks:
content.extend([{'type': 'text', 'text': f'<{i / 24:.3f} seconds>:'}, {'type': 'image', 'image': frames[i]}])
content.append({'type': 'text', 'text': 'Analyze the requested reference-to-performer assignments across these guide moments. '
'Keep reference identity separate from guide identity and motion. Do not invent unseen details. '
'Conclude with a short mapping summary. Request: ' + row['prompt']})
messages = [{'role': 'user', 'content': content}]
for thinking in (False, True):
dest = cache / f'{row["id"]}.student_{int(thinking)}.safetensors'
if dest.exists():
continue
inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True,
return_dict=True, return_tensors='pt', enable_thinking=thinking).to(model.device)
captured = []
def capture(_module, _inputs, output):
x = output[0] if isinstance(output, tuple) else output
assert x.ndim == 3 and x.shape[0] == 1 and x.shape[-1] == 4096
captured.append(x[0].detach().cpu().to(torch.bfloat16).contiguous())
handle = layer.register_forward_hook(capture)
try:
with torch.inference_mode():
generated = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False,
use_cache=True, return_dict_in_generate=False)
finally:
handle.remove()
assert captured, 'No decoder activations captured'
prefix = captured[0]; tail = torch.cat(captured[1:]) if len(captured) > 1 else prefix[:0]
text = processor.tokenizer.decode(generated[0, inputs['input_ids'].shape[1]:], skip_special_tokens=False)
completed = not thinking or '</think>' in text
meta = dict(input_key=row['input_key'], model=str(Path(args.model).resolve()), thinking=thinking,
thinking_completed=completed, prefix_tokens=len(prefix), generated_state_tokens=len(tail),
generated_tokens=generated.shape[1] - inputs['input_ids'].shape[1],
layer='last decoder block before final norm', visual_inputs=info,
extraction='Prefill activations plus generated-token activations; final sampled token may not be forwarded')
if not completed:
(cache / f'{row["id"]}.student_1.incomplete.json').write_text(json.dumps(meta, indent=2))
raise RuntimeError('Thinking did not finish within token budget. Increase --max-new-tokens; no valid thinking cache saved.')
save_file({'prefix': prefix, 'generated': tail}, str(dest), metadata={'experiment': json.dumps(meta)})
print(json.dumps(dict(id=row['id'], thinking=thinking, completed=completed,
prefix_tokens=len(prefix), generated_state_tokens=len(tail))), flush=True)
if __name__ == '__main__':
main()
@@ -0,0 +1,49 @@
"""GPU stage 1. Reuses the existing H3 encoder; does not load the H3 DiT/VAE."""
import argparse, importlib.util, json, sys
from pathlib import Path
from common import load_manifest, require_idle_gpu, visual_inputs
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--manifest', required=True); ap.add_argument('--cache', required=True)
ap.add_argument('--comfy', default='/data/comfyui')
ap.add_argument('--teacher', default='/data/models/minimax-h3/text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors')
args = ap.parse_args(); require_idle_gpu(); sys.argv = [sys.argv[0]]
import numpy as np
import torch
from safetensors.torch import save_file
sys.path.insert(0, args.comfy)
import comfy.options
comfy.options.enable_args_parsing()
import folder_paths, nodes
teacher = Path(args.teacher); folder_paths.add_model_folder_path('text_encoders', str(teacher.parent))
clip = nodes.CLIPLoader().load_clip(teacher.name, 'stable_diffusion')[0]
path = Path(__file__).resolve().parents[2] / 'minimax_h3_clip_video.py'
spec = importlib.util.spec_from_file_location('bridge_native_context', path)
module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
cache = Path(args.cache); cache.mkdir(parents=True, exist_ok=True)
for row in load_manifest(args.manifest):
dest = cache / (row['id'] + '.teacher.safetensors')
if dest.exists():
continue
refs, frames, picks, info = visual_inputs(row)
images = {f'ref_image_{i}': torch.from_numpy(np.asarray(im).copy()).float()[None] / 255 for i, im in enumerate(refs)}
guide = torch.from_numpy(np.stack([np.asarray(im) for im in frames]).copy()).float() / 255
positive = module.BFSMiniMaxH3ReferenceToVideoGuideContext.execute(
clip=clip, prompt=row['prompt'], width=info['width'], height=info['height'], length=info['length'],
ref_image_size='match', ref_images=images, guide_video_0=guide, clip_video_fps=2,
).args[0]
assert len(positive) == 1
embeds, extra = positive[0]; tags = extra['minimax_token_tags']
embeds = embeds[0] if embeds.ndim == 3 else embeds
tags = tags.reshape(-1)
assert embeds.shape == (len(tags), 5120) and torch.isfinite(embeds).all()
save_file({'teacher': embeds.cpu().to(torch.bfloat16).contiguous(), 'tags': tags.cpu().long().contiguous()},
str(dest), metadata={'input_key': row['input_key'], 'teacher': str(teacher), 'visual_inputs': json.dumps(info)})
(cache / (row['id'] + '.input.json')).write_text(json.dumps(row, indent=2, ensure_ascii=False))
print(json.dumps(dict(id=row['id'], teacher_shape=list(embeds.shape), split=row['split'])), flush=True)
if __name__ == '__main__':
main()
+53
View File
@@ -0,0 +1,53 @@
import json, subprocess
from pathlib import Path
def require_idle_gpu():
# No background polling/auto launch: fail before loading a model.
active = subprocess.check_output(['ps', '-eo', 'args'], text=True)
if any('audit_temporal_identity_binding.py' in line and not line.lstrip().startswith(('ps ', 'rg '))
for line in active.splitlines()):
raise RuntimeError('Dataset reviewer is active. No experiment GPU work was started.')
usage = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,noheader,nounits'], text=True)
if any(int(v) > 1024 for v in usage.split()):
raise RuntimeError('GPU is occupied. No experiment GPU work was started.')
def load_manifest(path):
rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
assert rows and len({r['id'] for r in rows}) == len(rows)
return rows
def visual_inputs(row, length=73, width=384, context_fps=2.0):
"""Identical 24-fps source interval and Picture order for teacher/student."""
import cv2
import numpy as np
from PIL import Image, ImageOps
cap = cv2.VideoCapture(row['guide']); fps = cap.get(cv2.CAP_PROP_FPS); count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if fps <= 0 or count <= 0:
raise ValueError('Unreadable guide ' + row['guide'])
n = min(length, int(count / fps * 24)); n -= (n - 5) % 17
if n < 5:
raise ValueError('Guide shorter than H3 grid')
frames = []
# Sequential decoding avoids failed random seeks on some original files.
wanted = [min(count - 1, round(i / 24 * fps)) for i in range(n)]; cursor = -1; bgr = None
for index in wanted:
while cursor < index:
ok, bgr = cap.read(); cursor += 1
if not ok:
raise ValueError(f'Guide decoding failed at source frame {cursor}')
frame = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
w = max(32, round(width / 32) * 32); h = max(32, round(frame.height / frame.width * w / 32) * 32)
frames.append(frame.resize((w, h), Image.Resampling.LANCZOS))
cap.release()
picks = sorted({round(i * 24 / context_fps) for i in range(int((n - 1) / 24 * context_fps) + 1)})
refs = []
for path in row['references']:
with Image.open(path) as raw:
im = raw.convert('RGB'); scale = min(1.0, (w * h / (im.width * im.height)) ** .5)
rw = max(32, round(im.width * scale / 32) * 32); rh = max(32, round(im.height * scale / 32) * 32)
refs.append(im.resize((rw, rh), Image.Resampling.LANCZOS))
return refs, frames, picks, dict(length=n, width=w, height=h, timestamps=[i / 24 for i in picks],
context_fps=context_fps, input_key=row['input_key'])
@@ -0,0 +1,51 @@
"""CPU-only input selection. These calibration inputs are NOT approved H3 training samples."""
import argparse, hashlib, json, random
from pathlib import Path
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--output', required=True)
parser.add_argument('--train', type=int, default=32)
parser.add_argument('--validation', type=int, default=8)
parser.add_argument('--test', type=int, default=8)
args = parser.parse_args()
rows = []
for family, root in [('genjutsu', Path('/data/datasets/openjutsu/train_mt_multi')),
('animate', Path('/data/datasets/animate_motion_transfer_v1/train'))]:
for caption in sorted((root / 'targets').glob('*.txt')):
refs = []
for i in range(1, 17):
files = sorted((root / f'refs_{i}').glob(caption.stem + '.*'))
if files:
refs.append(str(files[0]))
if not 1 <= len(refs) <= 4:
continue
guide = root / ('guides_original' if family == 'animate' else 'guides') / (caption.stem + '.mp4')
if not guide.is_file():
continue
rows.append(dict(id=family + '__' + caption.stem, guide=str(guide), references=refs,
prompt=caption.read_text().strip(), original_caption_sha256=hashlib.sha256(caption.read_bytes()).hexdigest(),
calibration_only=True, approved_for_h3_training=False))
random.Random(42).shuffle(rows)
selected = []; seen_assets = set()
# No exact reference-asset reuse across selected cases/splits.
for row in rows:
hashes = {hashlib.sha256(Path(p).read_bytes()).hexdigest() for p in row['references']}
if hashes & seen_assets:
continue
seen_assets.update(hashes); row['reference_sha256'] = sorted(hashes); selected.append(row)
if len(selected) >= args.train + args.validation + args.test:
break
assert len(selected) == args.train + args.validation + args.test, 'Insufficient disjoint reference assets'
for i, row in enumerate(selected):
row['split'] = 'train' if i < args.train else 'validation' if i < args.train + args.validation else 'test'
row['input_key'] = hashlib.sha256(json.dumps(row, sort_keys=True).encode()).hexdigest()
out = Path(args.output); out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in selected))
print(json.dumps(dict(output=str(out), total=len(selected), train=args.train, validation=args.validation, test=args.test,
purpose='Embedding reconstruction only; unfiltered input captions, not a curated identity-quality benchmark.')))
if __name__ == '__main__':
main()
@@ -0,0 +1,36 @@
"""Read-only CPU preflight: does not import torch, download weights or launch GPU jobs."""
import argparse, ast, json
from pathlib import Path
from common import load_manifest
def main():
ap = argparse.ArgumentParser(); ap.add_argument('--manifest', required=True)
ap.add_argument('--student', default='/data/models/Qwen3.5-9B')
ap.add_argument('--output', required=True); args = ap.parse_args()
rows = load_manifest(args.manifest)
missing = []
for row in rows:
for path in [row['guide'], *row['references']]:
if not Path(path).is_file():
missing.append(path)
for path in Path(__file__).parent.glob('*.py'):
ast.parse(path.read_text(), filename=str(path))
student = Path(args.student)
ready = (student / 'config.json').exists() and any(student.glob('*.safetensors'))
report = dict(phase='prepared_gpu_execution_pending', samples=len(rows),
splits={s: sum(r['split'] == s for r in rows) for s in ['train', 'validation', 'test']},
missing_input_files=missing, student_weights_found=ready, student_path=str(student),
syntax_checked=True, tensor_tests_run=False, gpu_tests_run=False, adapter_trained=False,
h3_quality_validated=False, active_dataset_review_untouched=True,
pending=['Provide full local Qwen3.5-9B Transformers weights in an isolated compatible environment.',
'After dataset GPU review exits, cache teacher then student activations sequentially.',
'Train shared bridge and evaluate held-out embedding reconstruction.',
'Render native baseline / thinking off / thinking on / generated-state ablation with fixed inputs.'])
if missing:
raise RuntimeError('Calibration input files missing: ' + str(missing))
Path(args.output).write_text(json.dumps(report, indent=2)); print(json.dumps(report, indent=2))
if __name__ == '__main__':
main()
@@ -0,0 +1,101 @@
"""GPU stage 3: one shared bridge, balanced thinking on/off, frozen cached encoders.
Embedding reconstruction is calibration only, NOT evidence that thinking helps H3.
"""
import argparse, json, random
from pathlib import Path
from common import load_manifest, require_idle_gpu
from bridge import BridgeConfig, HiddenStateBridge, alignment_loss
def main():
ap = argparse.ArgumentParser(); ap.add_argument('--manifest', required=True); ap.add_argument('--cache', required=True)
ap.add_argument('--output', required=True); ap.add_argument('--steps', type=int, default=500)
ap.add_argument('--eval-every', type=int, default=50); ap.add_argument('--lr', type=float, default=0.0001)
args = ap.parse_args(); require_idle_gpu()
import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file
torch.manual_seed(42); rng = random.Random(42)
rows = load_manifest(args.manifest); cache = Path(args.cache); out = Path(args.output)
if out.exists() and any(out.iterdir()):
raise RuntimeError('Refusing to overwrite an existing bridge run')
out.mkdir(parents=True, exist_ok=True)
for row in rows:
for suffix in ['teacher', 'student_0', 'student_1']:
path = cache / f'{row["id"]}.{suffix}.safetensors'
with safe_open(str(path), framework='pt') as f:
md = f.metadata() or {}
key = md.get('input_key') if suffix == 'teacher' else json.loads(md['experiment'])['input_key']
assert key == row['input_key'], 'Cache belongs to different inputs: ' + str(path)
train = [r for r in rows if r['split'] == 'train']; val = [r for r in rows if r['split'] == 'validation']
test = [r for r in rows if r['split'] == 'test']; assert train and val and test
def case(row, thinking):
teacher = load_file(str(cache / f'{row["id"]}.teacher.safetensors'))
student = load_file(str(cache / f'{row["id"]}.student_{thinking}.safetensors'))
return {k: v.cuda().float() for k, v in {**teacher, **student}.items()}
model = HiddenStateBridge(BridgeConfig()).cuda()
# Training-only per-channel scale; no held-out target features in model inputs.
squares = torch.zeros(5120, dtype=torch.float64); tokens = 0
for row in train:
t = load_file(str(cache / f'{row["id"]}.teacher.safetensors'))['teacher'].float()
squares += t.double().square().sum(0); tokens += len(t)
model.teacher_channel_scale.copy_((squares / tokens).sqrt().clamp_min(1e-3).float().cuda())
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=.01)
def predict(c, ablate=False):
generated = torch.zeros_like(c['generated']) if ablate else c['generated']
return model(c['prefix'], generated, c['tags'].long())
@torch.no_grad()
def measure(subset, export=False):
model.eval(); results = []
for row in subset:
for thinking in (0, 1):
c = case(row, thinking)
for ablate in ([False, True] if thinking and export else [False]):
pred = predict(c, ablate); t = c['teacher']
assert torch.isfinite(pred).all()
relative = float((pred - t).norm() / t.norm().clamp_min(1e-8))
cosine = float(torch.nn.functional.cosine_similarity(pred.flatten(), t.flatten(), dim=0))
results.append(dict(id=row['id'], split=row['split'], thinking=bool(thinking),
generated_states_zeroed=ablate, relative_error=relative, cosine=cosine,
loss=float(alignment_loss(pred, t, c['tags']))))
if export:
label = 'thinking_ablated' if ablate else f'thinking_{thinking}'
name = row['id'] + '.' + label + '.safetensors'
metadata = dict(experiment='h3_9b_hidden_state_bridge_v1', input_key=row['input_key'],
source_id=row['id'], thinking=str(bool(thinking)), generated_states_zeroed=str(ablate),
prompt=row['prompt'], validated_h3_quality='false', teacher_layout_required='true')
save_file({'conditioning': pred.detach().cpu().to(torch.bfloat16).contiguous(),
'teacher_baseline': t.cpu().to(torch.bfloat16).contiguous(),
'minimax_token_tags': c['tags'].long().cpu().contiguous()},
str(out / name), metadata=metadata)
return results
best = float('inf'); history = []
for step in range(1, args.steps + 1):
model.train(); row = rng.choice(train); thinking = step % 2
c = case(row, thinking); optimizer.zero_grad(set_to_none=True)
prediction = predict(c); loss = alignment_loss(prediction, c['teacher'], c['tags'])
assert torch.isfinite(loss)
loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); optimizer.step()
if step % args.eval_every == 0 or step == args.steps:
metrics = measure(val); score = sum(r['loss'] for r in metrics) / len(metrics)
entry = dict(step=step, train_loss=float(loss), validation_loss=score); history.append(entry)
print(json.dumps(entry), flush=True)
(out / 'history.json').write_text(json.dumps(history, indent=2))
if score < best:
best = score
save_file({k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()},
str(out / 'bridge.safetensors'), metadata={'config': json.dumps(model.config_dict()),
'selected_by': 'validation reconstruction loss',
'h3_quality_validated': 'false'})
model.load_state_dict(load_file(str(out / 'bridge.safetensors')))
metrics = measure(test, export=True)
report = dict(test_metrics=metrics, h3_quality_validated=False, thinking_benefit_proven=False,
next_step='Compare native teacher, thinking off, thinking on and generated-state ablation in H3 with fixed seed/guide/references.',
limitation='Teacher-supervised reconstruction can suppress thinking-specific information; lower embedding error is not identity/motion quality.')
(out / 'report.json').write_text(json.dumps(report, indent=2))
(out / 'manifest.jsonl').write_text(Path(args.manifest).read_text())
print(json.dumps(dict(phase='embedding_calibration_complete_gpu_video_validation_pending', output=str(out))), flush=True)
if __name__ == '__main__':
main()
+62
View File
@@ -0,0 +1,62 @@
"""Experimental cached-case embedding swap, preserving existing H3 latent conditioning."""
import os
import torch
import folder_paths
from safetensors import safe_open
from safetensors.torch import load_file
folder_paths.add_model_folder_path('h3_embedding_bridge', os.path.join(folder_paths.models_dir, 'h3_embedding_bridge'))
class BFSH3ExperimentalEmbeddingBridge:
@classmethod
def INPUT_TYPES(cls):
return {'required': {
'positive': ('CONDITIONING',),
'embedding_file': (folder_paths.get_filename_list('h3_embedding_bridge') or ['(no calibrated export)'],),
'baseline_tolerance': ('FLOAT', {'default': .02, 'min': 0.0, 'max': .05, 'step': .001,
'tooltip': 'Reject cached embeddings for a different input/encoder. Use the matching calibration case, prompt, media interval and encoder. Do not raise this to bypass mismatched inputs.'}),
}}
RETURN_TYPES = ('CONDITIONING', 'STRING')
RETURN_NAMES = ('positive', 'experiment_info')
FUNCTION = 'apply'
CATEGORY = 'BFS/experimental'
DESCRIPTION = ('Replace only text/vision embeddings with a calibrated 9B hidden-state bridge export for the SAME cached case. '
'Preserves native reference latents and aligned guides. Requires trained bridge exports; this is not a general encoder replacement.')
def apply(self, positive, embedding_file, baseline_tolerance=.02):
if len(positive) != 1:
raise ValueError('Experimental bridge supports one native H3 conditioning entry')
path = folder_paths.get_full_path_or_raise('h3_embedding_bridge', embedding_file)
with safe_open(path, framework='pt') as f:
md = f.metadata() or {}
if md.get('experiment') != 'h3_9b_hidden_state_bridge_v1':
raise ValueError('Select a held-out conditioning export, not bridge weights or another LoRA')
data = load_file(path); base = positive[0][0]; source = base[0] if base.ndim == 3 else base
expected = data['teacher_baseline'].float(); actual = source.detach().cpu().float()
if actual.shape != expected.shape:
raise ValueError('Cached case does not match current prompt/media token layout')
relative = float((actual - expected).norm() / expected.norm().clamp_min(1e-8))
if relative > baseline_tolerance:
raise ValueError(f'Cached case mismatch: baseline relative error {relative:.5f}; use exact matching inputs')
replacement = data['conditioning']; tags = data['minimax_token_tags']
if replacement.shape != expected.shape or len(tags) != len(replacement) or not torch.isfinite(replacement).all():
raise ValueError('Invalid bridge conditioning export')
current_tags = positive[0][1].get('minimax_token_tags')
if current_tags is None or not torch.equal(current_tags.detach().cpu().reshape(-1).long(), tags.reshape(-1).long()):
raise ValueError('Cached case has different modality tags')
values = dict(positive[0][1])
values['minimax_token_tags'] = tags.to(current_tags.device)
replacement = replacement.to(device=base.device, dtype=base.dtype)
if base.ndim == 3:
replacement = replacement[None]
info = (f"EXPERIMENTAL {md['source_id']}: thinking={md['thinking']}, "
f"generated states zeroed={md['generated_states_zeroed']}; "
f"baseline error={relative:.5f}; H3 quality not validated")
print('[BFS H3 embedding bridge] ' + info, flush=True)
return [[replacement, values]], info
NODE_CLASS_MAPPINGS = {'BFSH3ExperimentalEmbeddingBridge': BFSH3ExperimentalEmbeddingBridge}
NODE_DISPLAY_NAME_MAPPINGS = {'BFSH3ExperimentalEmbeddingBridge': 'BFS H3 Experimental 9B Embedding Bridge (Cached Case)'}