| |
| """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: |
| |
| 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() |
|
|
