File size: 3,462 Bytes
54daaf9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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() |