Blame for infer.py

"""Single-GPU Kelix-AEDIT inference: an image or local VQ indices -> PNG."""
import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
# Support running the demo or a downloaded HF script inside a Kelix checkout.
for parent in (ROOT, *ROOT.parents):
    if (parent / 'muse' / '__init__.py').is_file():
        sys.path.insert(0, str(parent))
        break


def parse_args():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--ckpt', default=str(ROOT) if (ROOT / 'config.json').is_file() else 'OpenOneRec/Kelix-AEDIT', help='Local DiT directory or HF repo ID')
    p.add_argument('--tokenizer', default='OpenOneRec/Kelix-SFT', help='Kelix-SFT repo/directory or a standalone image tokenizer directory')
    p.add_argument('--vae', default='Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers/vae', help='Local DC-AE directory or HF repo/subfolder')
    inputs = p.add_mutually_exclusive_group(required=True)
    inputs.add_argument('--image', help='Image to encode and reconstruct')
    inputs.add_argument('--tokens', help='Safely loaded .pt LongTensor [H*W,8], local codebook indices')
    p.add_argument('--grid', type=int, default=18, help='Condition grid width/height; 18 means 324 tokens')
    p.add_argument('--output', default='output.png')
    p.add_argument('--image-size', type=int, default=512)
    p.add_argument('--steps', type=int, default=20)
    p.add_argument('--cfg', type=float, default=4.5)
    p.add_argument('--flow-shift', type=float, default=3.0)
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--device', default='cuda:0')
    p.add_argument('--dtype', choices=['bfloat16', 'float16', 'float32'], default='bfloat16')
    return p.parse_args()


def positions(h, w, device):
    import torch
    ids = torch.arange(h * w, device=device)
    return {'height': ids // w, 'width': ids % w}


def resolve_directory(source, allow_patterns=None):
    path = Path(source)
    if path.is_dir():
        return path
    from huggingface_hub import snapshot_download
    parts = str(source).split('/')
    repo_id = '/'.join(parts[:2])
    subfolder = '/'.join(parts[2:])
    patterns = [f'{subfolder}/*'] if subfolder else allow_patterns
    root = Path(snapshot_download(repo_id, allow_patterns=patterns))
    return root / subfolder if subfolder else root


def load_image_tokenizer(source, device, dtype):
    """Load the visual tokenizer without instantiating or downloading the LLM."""
    import torch
    from accelerate import init_empty_weights
    from safetensors import safe_open
    from muse.config.model_config import KeyeTokenizerConfig
    from muse.models.keye_tokenizer import KeyeImageTokenizer

    folder = resolve_directory(source, ['*.json', '*.txt', '*.py', '*.jinja'])
    config = json.loads((folder / 'config.json').read_text())
    if 'tokenizer_config' not in config:
        tokenizer = KeyeImageTokenizer.from_pretrained(
            str(folder), strict=True, fusion_type='mean'
        ).to(device=device, dtype=dtype).eval()
        return tokenizer, folder

    tokenizer_config = KeyeTokenizerConfig(**config['tokenizer_config'])
    tokenizer_config.fusion_type = 'mean'
    with init_empty_weights():
        tokenizer = KeyeImageTokenizer(tokenizer_config)
    index = json.loads((folder / 'model.safetensors.index.json').read_text())
    prefix = 'visual_tokenizer.'
    selected = {k: v for k, v in index['weight_map'].items() if k.startswith(prefix)}
    if not selected:
        raise ValueError('No visual_tokenizer weights found in the checkpoint index')
    state = {}
    for shard in sorted(set(selected.values())):
        path = folder / shard
        if not path.is_file():
            if Path(source).is_dir():
                raise FileNotFoundError(f'Missing tokenizer shard: {shard}')
            from huggingface_hub import hf_hub_download
            path = Path(hf_hub_download(source, shard, revision=folder.name))
        with safe_open(str(path), framework='pt', device='cpu') as weights:
            for key, filename in selected.items():
                if filename == shard:
                    state[key[len(prefix):]] = weights.get_tensor(key)
    tokenizer.load_state_dict(state, strict=True, assign=True)
    tokenizer = tokenizer.to(device=device, dtype=dtype).eval()
    return tokenizer, folder


def main(args):
    import torch
    import torch.nn.functional as F
    from PIL import Image
    from diffusers import AutoencoderDC, FlowMatchEulerDiscreteScheduler
    from transformers import AutoProcessor
    from muse.models.base import Model
    from muse.training.common import set_default_dtype

    if args.grid < 1 or args.steps < 1 or args.image_size < 32 or args.image_size % 32:
        raise ValueError('grid and steps must be positive; image-size must be a positive multiple of 32')
    ckpt = resolve_directory(args.ckpt, ['config.json', '*.safetensors', '*.index.json'])
    dit_dir = ckpt / 'dit' if (ckpt / 'dit').is_dir() else ckpt
    device, dtype = torch.device(args.device), getattr(torch, args.dtype)
    if device.type == 'cuda' and not torch.cuda.is_available():
        raise RuntimeError('CUDA is required for the FlashAttention vision encoder')
    torch.manual_seed(args.seed)
    with torch.inference_mode(), set_default_dtype(dtype):
        print('Loading DiT (strict=True)...', flush=True)
        model = Model.from_pretrained(str(dit_dir), strict=True).to(device=device, dtype=dtype).eval()
        print('Loading image tokenizer (strict=True)...', flush=True)
        tokenizer, tokenizer_dir = load_image_tokenizer(args.tokenizer, device, dtype)
        vae_dir = resolve_directory(args.vae, ['config.json', '*.safetensors'])
        vae = AutoencoderDC.from_pretrained(str(vae_dir), torch_dtype=dtype).to(device).eval()
        length = args.grid ** 2
        if args.tokens:
            indices = torch.load(args.tokens, map_location='cpu', weights_only=True)
            if not isinstance(indices, torch.Tensor) or tuple(indices.shape) != (length, tokenizer.n_q_tokens):
                raise ValueError(f'Expected integer tensor [{length}, {tokenizer.n_q_tokens}]')
            if indices.dtype not in (torch.int32, torch.int64):
                raise ValueError('Token indices must have integer dtype')
            limit = tokenizer.config.codebook_size // tokenizer.config.split_voc
            if indices.min() < 0 or indices.max() >= limit:
                raise ValueError(f'Expected per-codebook indices in [0, {limit}) (not AR vocabulary IDs)')
            cond, mask, seq_len = tokenizer.tokenize_indices(
                indices.to(device=device, dtype=torch.long), max_pad_to=length,
                lengths=[length], return_attention_mask=True,
            )
        else:
            processor = AutoProcessor.from_pretrained(
                str(tokenizer_dir), trust_remote_code=True, use_fast=False,
            )
            # Match the original inference's square resize, then use an exact grid.
            # A merged Keye token covers 2*14 pixels; 18*28=504 for 324 tokens.
            image = Image.open(args.image).convert('RGB').resize(
                (args.image_size, args.image_size), Image.Resampling.LANCZOS,
            ).resize((args.grid * 28, args.grid * 28), Image.Resampling.BICUBIC)
            inputs = processor(text=[''], images=[image], return_tensors='pt')
            actual_grid = inputs['image_grid_thw'][0, 1:] // 2
            if actual_grid.tolist() != [args.grid, args.grid]:
                raise ValueError(f'Unexpected processor grid: {actual_grid.tolist()}')
            encoded = tokenizer(
                pixel_values=inputs['pixel_values'].to(device=device, dtype=dtype),
                image_grid_thw=inputs['image_grid_thw'].to(device),
            )
            features = encoded.get('token_embeds')
            if features is None:
                # Older Kelix releases compute fused features without returning them.
                features = torch.stack([
                    projector(z_q) for projector, z_q in
                    zip(tokenizer.up_projectors, encoded['z_q'])
                ], dim=1).mean(dim=1)
            if features.shape[0] != length:
                raise ValueError(f'Expected {length} image tokens, got {features.shape[0]}')
            cond, seq_len = features[None], length
            mask = torch.ones((1, 1, 1, length), device=device)
        if cond.shape[-1] != model.config.caption_channels:
            raise ValueError('Tokenizer condition dimension does not match the DiT config')
        null = model.y_embedder.y_embedding[:seq_len].to(device=device, dtype=dtype)
        valid = null.shape[0]
        null = F.pad(null, (0, 0, 0, seq_len - valid))[None]
        null_mask = torch.zeros_like(mask)
        null_mask[..., :valid] = 1
        cond_cfg = torch.cat([null, cond])
        mask_cfg = torch.cat([null_mask, mask])
        latent_size = args.image_size // 32
        generator = torch.Generator(device=device).manual_seed(args.seed)
        latents = torch.randn((1, model.config.in_channels, latent_size, latent_size),
                              device=device, dtype=dtype, generator=generator)
        scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
        scheduler.set_timesteps(args.steps, device=device)
        x_pos = positions(latent_size, latent_size, device)
        y_pos = positions(args.grid, args.grid, device)
        print(f'Sampling {args.steps} steps...', flush=True)
        for t in scheduler.timesteps:
            pred = model.forward_with_dpmsolver(
                torch.cat([latents, latents]), t.expand(2), cond_cfg, mask=mask_cfg,
                x_input_pos=x_pos, cond_input_pos=y_pos,
            )
            uncond, conditional = pred.chunk(2)
            pred = uncond + args.cfg * (conditional - uncond)
            latents = scheduler.step(pred, t, latents, return_dict=False)[0]
        image = (vae.decode(latents / vae.config.scaling_factor).sample[0] / 2 + 0.5).clamp(0, 1)
        if not torch.isfinite(image).all():
            raise RuntimeError('Non-finite pixels encountered')
        pixels = (image.float().permute(1, 2, 0).cpu().numpy() * 255).round().astype('uint8')
        output = Path(args.output)
        output.parent.mkdir(parents=True, exist_ok=True)
        Image.fromarray(pixels).save(output)
        print(f'Saved {output.resolve()}')


if __name__ == '__main__':
    main(parse_args())