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