"""Component-aware vLLM-Omni inference for the full-encoder checkpoint."""
import argparse
import json
import os
import re
import tempfile
from pathlib import Path

os.environ.setdefault('CUDA_LAUNCH_BLOCKING', '1')


def component_config(model, component):
    from safetensors import safe_open

    folder = model / component
    cfg = json.loads((folder / 'quantization_config.json').read_text())
    cfg.update(block_name_to_quantize=None, extra_config={})
    for shard in sorted(folder.glob('*.safetensors')):
        with safe_open(shard, framework='pt', device='cpu') as handle:
            keys = set(handle.keys())
            for key in keys:
                if not key.endswith('.weight'):
                    continue
                layer = key.removesuffix('.weight')
                if layer + '.weight_scale' not in keys and len(handle.get_slice(key).get_shape()) == 2:
                    cfg['extra_config']['^' + re.escape(component + '.' + layer) + '$'] = {
                        'bits': 16, 'data_type': 'float', 'act_bits': 16, 'act_data_type': 'float',
                    }
    return cfg


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument('--prompt', default='A red colored car.')
    parser.add_argument('--output', type=Path, default=Path('output.png'))
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    from vllm_omni.entrypoints.omni import Omni
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams

    model = args.model.resolve()
    configs = {name: component_config(model, name) for name in ('transformer', 'text_encoder')}
    configs['vae'] = None
    # A fresh view avoids stale links and duplicate safetensors index files.
    with tempfile.TemporaryDirectory(prefix='flux2-runtime-') as directory:
        view = Path(directory)
        for name in ('transformer', 'text_encoder'):
            target = view / name
            target.mkdir()
            new_index = (model / name / 'model.safetensors.index.json').exists()
            for source in (model / name).iterdir():
                if new_index and source.name == 'diffusion_pytorch_model.safetensors.index.json':
                    continue
                (target / source.name).symlink_to(source, target_is_directory=source.is_dir())
        for name in ('model_index.json', 'vae', 'scheduler', 'tokenizer'):
            (view / name).symlink_to(model / name, target_is_directory=(model / name).is_dir())
        omni = Omni(model=str(view), quantization_config=configs,
                    init_timeout=3600, stage_init_timeout=3600)
        params = OmniDiffusionSamplingParams(height=1024, width=1024, seed=args.seed,
                                            guidance_scale=4.0, num_inference_steps=50,
                                            num_outputs_per_prompt=1)
        outputs = omni.generate({'prompt': args.prompt, 'modalities': ['image']},
                                sampling_params_list=[params])
        for output in outputs:
            if getattr(output, 'images', None):
                args.output.parent.mkdir(parents=True, exist_ok=True)
                output.images[0].save(args.output)
                print(args.output)
                break
        else:
            raise RuntimeError('No image returned')


if __name__ == '__main__':
    main()