| import hashlib |
| import json |
| from pathlib import Path |
|
|
|
|
| BASE_MODEL = 'Qwen/Qwen-Image-2.1' |
| VARIANTS = tuple(f'{steps}step-{precision}' for steps in (2, 4, 8) for precision in ('bf16', 'fp8')) |
|
|
|
|
| def resolve_model(model, variant, source): |
| if model: |
| folder = Path(model).expanduser().resolve() |
| elif source == 'huggingface': |
| from huggingface_hub import snapshot_download |
|
|
| folder = Path(snapshot_download( |
| 'Wuli-art/Qwen-Image-2.1-Turbo', allow_patterns=[f'{variant}/*'], |
| )) / variant |
| else: |
| from modelscope import snapshot_download |
|
|
| folder = Path(snapshot_download( |
| 'Wuli-Art/Qwen-Image-2.1-Turbo', allow_patterns=[f'{variant}/*'], |
| )) / variant |
| if not (folder / 'sampling_config.json').is_file(): |
| raise ValueError('Select a DiT directory containing sampling_config.json') |
| return folder |
|
|
|
|
| def read_config(folder): |
| sampling = json.loads((folder / 'sampling_config.json').read_text()) |
| transformer = json.loads((folder / 'config.json').read_text()) |
| sigmas = sampling['sample_sigmas'] |
| if len(sigmas) not in (2, 4, 8) or sigmas[0] != 1.0: |
| raise ValueError('Unexpected saved sampling schedule') |
| if any(first <= second for first, second in zip(sigmas, sigmas[1:])): |
| raise ValueError('Sampling sigmas must strictly decrease') |
| return sampling, transformer |
|
|
|
|
| def resolve_base(model=BASE_MODEL, source='huggingface'): |
| path = Path(model).expanduser() |
| if path.is_dir(): |
| return path.resolve() |
| patterns = ['model_index.json', 'processor/*', 'scheduler/*', 'text_encoder/*', 'vae/*'] |
| if source == 'huggingface': |
| from huggingface_hub import snapshot_download |
|
|
| return Path(snapshot_download(model, allow_patterns=patterns)) |
| from modelscope import snapshot_download |
|
|
| return Path(snapshot_download(model, allow_patterns=patterns)) |
|
|
|
|
| def prepare_pipeline_directory(dit_folder, base_folder): |
| sampling, transformer = read_config(dit_folder) |
| key = hashlib.sha256((str(dit_folder.resolve()) + str(base_folder.resolve())).encode()).hexdigest()[:20] |
| destination = Path.home() / '.cache' / 'wuli-qwen-image21' / key |
| destination.mkdir(parents=True, exist_ok=True) |
| for component in ['processor', 'text_encoder', 'vae']: |
| target = base_folder / component |
| if not target.is_dir(): |
| raise ValueError(f'Missing base component: {component}') |
| link = destination / component |
| if not link.exists(): |
| link.symlink_to(target.resolve(), target_is_directory=True) |
| link = destination / 'transformer' |
| if not link.exists(): |
| link.symlink_to(dit_folder.resolve(), target_is_directory=True) |
| scheduler = destination / 'scheduler' |
| scheduler.mkdir(exist_ok=True) |
| scheduler_config = json.loads((dit_folder / 'scheduler_config.json').read_text()) |
| (scheduler / 'scheduler_config.json').write_text(json.dumps(scheduler_config, indent=2) + '\n') |
| pipeline = json.loads((base_folder / 'model_index.json').read_text()) |
| pipeline['sample_sigmas'] = sampling['sample_sigmas'] |
| (destination / 'model_index.json').write_text(json.dumps(pipeline, indent=2) + '\n') |
| return destination |
|
|
