#!/usr/bin/env python3 """Training-free NF editing using a fixed source-image score reference. Update the scaled VAE latent with the target score minus the cached source score. Both branches query the full conditional log density with the same fixed Gaussian noise, at the noise scale from the model configuration. """ import argparse from dataclasses import asdict, dataclass from io import BytesIO from itertools import chain, islice import json import math import os from pathlib import Path import sys import time import torch from torch.utils.checkpoint import checkpoint ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) DEFAULT_DATASET = Path('/mnt/bn/jiangzhongtao/lingmin/dataset/PIE_Bench_pp') @dataclass class EditingConfig: noise_std: float = 0.3 num_steps: int = 50 step_size: float = 1e-3 noise_batch_size: int = 1 gradient_checkpointing: bool = True def conditional_density(model, latent, condition): z, _, _, logdets = model(latent, condition) flat_z = z.flatten(1).double() prior_nll = 0.5 * flat_z.square().sum(1) prior_nll = prior_nll + 0.5 * flat_z.shape[1] * math.log(2 * math.pi) log_abs_det = torch.zeros_like(prior_nll) for logdet in logdets: log_abs_det = log_abs_det + logdet.flatten(1).double().sum(1) return torch.stack((log_abs_det - prior_nll, prior_nll, log_abs_det), dim=1) @torch.enable_grad() def conditional_score(model, latent, condition, gradient_checkpointing=True): query = latent.detach().requires_grad_(True) if gradient_checkpointing: # The model's built-in checkpointing is disabled in eval mode. density = checkpoint(lambda u, c: conditional_density(model, u, c), query, condition.detach(), use_reentrant=False) else: density = conditional_density(model, query, condition.detach()) score, = torch.autograd.grad(density[:, 0].sum(), query) if not torch.isfinite(density).all() or not torch.isfinite(score).all(): raise FloatingPointError('Non-finite NF density or score; no update was applied') return score.detach(), density.detach() @torch.no_grad() def mean_score(model, latent, condition, noise, config): score_sum = torch.zeros_like(latent) density_sum = torch.zeros(3, device=latent.device, dtype=torch.float64) for chunk in noise.split(config.noise_batch_size): query = latent + config.noise_std * chunk repeated_condition = condition.expand(chunk.shape[0], *condition.shape[1:]) score, density = conditional_score(model, query, repeated_condition, config.gradient_checkpointing) score_sum += score.sum(0, keepdim=True) density_sum += density.sum(0) return score_sum / noise.shape[0], density_sum / noise.shape[0] @torch.no_grad() def edit_latent(model, source_latent, source_condition, target_condition, noise, config, verbose=True): if source_latent.shape[0] != 1: raise ValueError('Edit one image at a time; noise_batch_size batches noise queries only') if noise.shape[0] < 1 or noise.shape[1:] != source_latent.shape[1:]: raise ValueError('noise must have shape [M, *source_latent.shape[1:]], with M >= 1') source = source_latent.detach().float().clone() current = source.clone() noise = noise.detach() source_score, source_density = mean_score(model, source, source_condition, noise, config) trajectory = [] for step in range(config.num_steps): target_score, target_density = mean_score(model, current, target_condition, noise, config) direction = target_score - source_score objective = (target_density[0] - source_density[0] - (source_score.double() * (current.double() - source.double())).sum()) updated = current + config.step_size * direction if not torch.isfinite(updated).all(): raise FloatingPointError(f'Non-finite latent at update {step + 1}; reduce --step-size') state = { 'step': step + 1, 'target_log_density_before_update': target_density[0].item(), 'target_prior_nll_before_update': target_density[1].item(), 'target_log_abs_det_before_update': target_density[2].item(), 'objective_before_update': objective.item(), 'target_score_rms': target_score.double().square().mean().sqrt().item(), 'direction_rms': direction.double().square().mean().sqrt().item(), 'update_rms': (updated.double() - current.double()).square().mean().sqrt().item(), 'displacement_rms_after_update': (updated.double() - source.double()).square().mean().sqrt().item(), } trajectory.append(state) current = updated.detach() if verbose: print(f" step={step + 1}/{config.num_steps} " f"Phi_before={state['objective_before_update']:.6g} " f"direction_rms={state['direction_rms']:.6g} " f"update_rms={state['update_rms']:.6g} " f"displacement_rms={state['displacement_rms_after_update']:.6g}", flush=True) result = { 'completed_steps': config.num_steps, 'source_log_density': source_density[0].item(), 'source_prior_nll': source_density[1].item(), 'source_log_abs_det': source_density[2].item(), 'source_score_rms': source_score.double().square().mean().sqrt().item(), 'final_displacement_rms': (current.double() - source.double()).square().mean().sqrt().item(), 'nf_score_calls': (config.num_steps + 1) * math.ceil(noise.shape[0] / config.noise_batch_size), 'forward_calls_exclude_checkpoint_recomputation': True, 'trajectory_density_timing': 'before each update; no extra final density query', } return current, source_score, result, trajectory def parse_args(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--image-path', type=Path, help='Omit to use PIE-Bench++') parser.add_argument('--source-prompt') parser.add_argument('--target-prompt') parser.add_argument('--dataset', type=Path, default=DEFAULT_DATASET) parser.add_argument('--category', action='append', help='Repeatable PIE category filter') parser.add_argument('--image-id', type=int, action='append', help='Repeatable PIE image ID filter') parser.add_argument('--start-index', type=int, default=0) parser.add_argument('--max-samples', '--max-sample', type=int, default=0, help='0: all selected images') parser.add_argument('--output-dir', type=Path, default=ROOT / 'result/nf_reference_score', help='Must not already exist') parser.add_argument('--model-config-path', type=Path, default=ROOT / 'configs/starflow_3B_t2i_256x256.yaml') parser.add_argument('--checkpoint-path', type=Path, default=ROOT / 'ckpts/starflow_3B_t2i_256x256.pth') parser.add_argument('--seed', type=int, default=999) parser.add_argument('--vae-latent', choices=('mean', 'sample'), default='mean') parser.add_argument('--num-noise-samples', type=int, default=4, help='M; fixed for all iterations') parser.add_argument('--noise-batch-size', type=int, default=EditingConfig.noise_batch_size, help='Noise queries per NF call; reduce for lower peak memory') parser.add_argument('--num-steps', '--max-steps', type=int, default=EditingConfig.num_steps, help='K') parser.add_argument('--step-size', type=float, default=EditingConfig.step_size, help='eta; summed-density score') parser.add_argument('--no-gradient-checkpointing', dest='gradient_checkpointing', action='store_false') cli = parser.parse_args() cli.precision = 'fp32' if cli.image_path is not None: if not cli.image_path.is_file(): parser.error('--image-path must be an existing file') if not cli.source_prompt or not cli.source_prompt.strip() or not cli.target_prompt or not cli.target_prompt.strip(): parser.error('Single-image mode requires nonempty --source-prompt and --target-prompt') if cli.category or cli.image_id or cli.start_index or cli.max_samples: parser.error('PIE selection options cannot be used with --image-path') elif cli.source_prompt is not None or cli.target_prompt is not None: parser.error('Prompt overrides require --image-path') if cli.num_noise_samples < 1 or cli.noise_batch_size < 1: parser.error('num-noise-samples and noise-batch-size must be positive') if min(cli.num_steps, cli.start_index, cli.max_samples) < 0: parser.error('num-steps, start-index and max-samples must be nonnegative') if cli.image_id and min(cli.image_id) < 0: parser.error('image-id must be nonnegative') if not math.isfinite(cli.step_size) or cli.step_size <= 0: parser.error('step-size must be finite and positive') if int(os.environ.get('WORLD_SIZE', '1')) != 1: parser.error('Use one process/GPU') for name in ('image_path', 'dataset', 'output_dir', 'model_config_path', 'checkpoint_path'): value = getattr(cli, name) if value is not None: setattr(cli, name, value.resolve()) return cli def iter_inputs(cli): if cli.image_path is not None: yield {'category': 'single', 'id': cli.image_path.stem, 'image_bytes': cli.image_path.read_bytes(), 'source_prompt': cli.source_prompt.strip(), 'target_prompt': cli.target_prompt.strip()} return import pyarrow.parquet as pq paths = sorted(cli.dataset.glob('*/*.parquet')) categories = {path.parent.name for path in paths} if cli.category: unknown = set(cli.category) - categories if unknown: raise ValueError(f'Unknown categories: {sorted(unknown)}; available: {sorted(categories)}') paths = [path for path in paths if path.parent.name in cli.category] if not paths: raise FileNotFoundError(f'No PIE-Bench++ Parquet files under {cli.dataset}') index = 0 requested, found = set(cli.image_id or []), set() for path in paths: parquet = pq.ParquetFile(path) for batch in parquet.iter_batches(batch_size=64, columns=['image', 'id', 'source_prompt', 'target_prompt']): for row in batch.to_pylist(): index += 1 if index <= cli.start_index or (requested and int(row['id']) not in requested): continue found.add(int(row['id'])) yield {'category': path.parent.name, 'id': str(row['id']), 'image_bytes': row['image']['bytes'], 'source_prompt': row['source_prompt'].strip(), 'target_prompt': row['target_prompt'].replace('[', '').replace(']', '').strip()} if requested - found: raise ValueError(f'IDs missing after category/start-index filters: {sorted(requested - found)}') def preprocess_image(image_bytes, image_size): from PIL import Image, ImageOps import torchvision.transforms as transforms with Image.open(BytesIO(image_bytes)) as raw: image = ImageOps.exif_transpose(raw).convert('RGB') scale = max(image_size / image.height, image_size / image.width) transform = transforms.Compose([ transforms.Resize((int(image.height * scale), int(image.width * scale))), transforms.CenterCrop((image_size, image_size)), transforms.ToTensor(), transforms.Normalize([0.5] * 3, [0.5] * 3), ]) return transform(image).unsqueeze(0) def save_image(tensor, path): from PIL import Image if not torch.isfinite(tensor).all(): raise FloatingPointError(f'Non-finite decoded image: {path}') path.parent.mkdir(parents=True, exist_ok=True) pixels = (tensor[0].detach().float().cpu().clamp(-1, 1) + 1) * 127.5 Image.fromarray(pixels.round().byte().permute(1, 2, 0).numpy()).save(path) def main(): cli = parse_args() if cli.output_dir.exists(): raise FileExistsError(f'Choose a new --output-dir: {cli.output_dir}') inputs = iter_inputs(cli) if cli.image_id: inputs = iter(list(inputs)) if cli.max_samples: inputs = islice(inputs, cli.max_samples) first = next(inputs, None) if first is None: raise ValueError('No images selected') if not torch.cuda.is_available(): raise RuntimeError('STARFlow 3B editing requires a CUDA GPU') from pie_bench_nf_components import load_components from utils import encode_text os.chdir(ROOT) args, model, vae, tokenizer, text_encoder, dist, device = load_components(cli) for component in (model, vae, text_encoder): component.eval().requires_grad_(False) config = EditingConfig(noise_std=float(args.noise_std), num_steps=cli.num_steps, step_size=cli.step_size, noise_batch_size=cli.noise_batch_size, gradient_checkpointing=cli.gradient_checkpointing) if args.noise_type != 'gaussian' or not math.isfinite(config.noise_std) or config.noise_std < 0: raise ValueError('The model configuration must use a finite, nonnegative Gaussian noise_std') cli.output_dir.mkdir(parents=True, exist_ok=False) run_config = {key: str(value) if isinstance(value, Path) else value for key, value in vars(cli).items()} run_config.update(editing=asdict(config), density_reduction='sum', density_accumulation='fp64', score_precision='fp32', guidance=0, fixed_source_reference=True, fixed_noise=True, torch_version=torch.__version__, gpu=torch.cuda.get_device_name(device), density_space='scaled VAE latent', finetuned_vae='none', preprocessing='resize shortest side then center crop to model image size') (cli.output_dir / 'run_config.json').write_text(json.dumps(run_config, indent=2, ensure_ascii=False) + '\n') previous_tf32 = torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False reports = [] try: for item in chain([first], inputs): started = time.perf_counter() category, image_id = item['category'], item['id'] if not item['source_prompt'] or not item['target_prompt']: raise ValueError(f'Missing prompt: {category}/{image_id}') print(f'Processing {category}/{image_id}; sigma={config.noise_std}; M={cli.num_noise_samples}', flush=True) image = preprocess_image(item['image_bytes'], args.img_size * vae.downsample_factor).to(device) item_seed = cli.seed + sum(ord(c) for c in category) + ( int(image_id) if image_id.isdigit() else sum(ord(c) for c in image_id)) generator = torch.Generator(device=device).manual_seed(item_seed) with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16): posterior = vae.vae.encode(image).latent_dist encoded = posterior.sample(generator=generator) if cli.vae_latent == 'sample' else posterior.mode() source_latent = (encoded * vae.scaling_factor).detach().float() conditions = encode_text( text_encoder, tokenizer, [item['source_prompt'], item['target_prompt']], args.txt_size, device, noise_std=config.noise_std if args.cond_noise_level else None).float() vae_only = vae.decode(source_latent) noise = torch.randn((cli.num_noise_samples, *source_latent.shape[1:]), generator=generator, device=device, dtype=torch.float32) with torch.autocast('cuda', enabled=False): edited_latent, source_score, result, trajectory = edit_latent( model, source_latent, conditions[:1], conditions[1:], noise, config) with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16): edited = vae.decode(edited_latent) stem = Path(category) / image_id for kind, pixels in (('source', image), ('vae_only', vae_only), ('editing', edited)): save_image(pixels, cli.output_dir / kind / stem.with_name(stem.name + '.png')) for folder in ('metadata', 'trajectory', 'latents'): (cli.output_dir / folder / stem.parent).mkdir(parents=True, exist_ok=True) result.update(category=category, id=image_id, source_prompt=item['source_prompt'], target_prompt=item['target_prompt'], item_seed=item_seed, noise_std=config.noise_std, num_noise_samples=cli.num_noise_samples, elapsed_seconds=time.perf_counter() - started) (cli.output_dir / 'metadata' / stem.with_name(stem.name + '.json')).write_text( json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False) + '\n') with (cli.output_dir / 'trajectory' / stem.with_name(stem.name + '.jsonl')).open('x') as handle: for state in trajectory: handle.write(json.dumps(state, allow_nan=False) + '\n') torch.save({'source_latent': source_latent.cpu(), 'edited_latent': edited_latent.cpu(), 'noise': noise.cpu(), 'source_score': source_score.cpu()}, cli.output_dir / 'latents' / stem.with_name(stem.name + '.pt')) reports.append(result) finally: torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = previous_tf32 if dist.distributed: torch.distributed.destroy_process_group() summary = {'count': len(reports), 'images': reports} (cli.output_dir / 'summary.json').write_text( json.dumps(summary, indent=2, ensure_ascii=False, allow_nan=False) + '\n') print(f'Finished {len(reports)} images; output={cli.output_dir}', flush=True) if __name__ == '__main__': main()