| """Validate pipeline structure and bind an evaluation to exact files.""" |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| from scripts.integrity import sha256 |
|
|
| REVISION = 'b3179ad355be050328e483a9dfdd9e60cd62adfa' |
|
|
|
|
| def _quantization(config): |
| raw = config.get('quantization_config') or {} |
| method = str(raw.get('quant_method', '')).lower() |
| dtype = str(raw.get('weights_dtype', '')).lower() |
| return method == 'sdnq' and dtype == 'uint4' and not raw.get('use_quantized_matmul') |
|
|
|
|
| def validate_model_structure(root): |
| root = Path(root) |
| required = ['model_index.json', 'processor/tokenizer.json', |
| 'processor/tokenizer_config.json', 'processor/preprocessor_config.json', |
| 'scheduler/scheduler_config.json'] |
| required += [f'{component}/config.json' for component in ('transformer', 'text_encoder', 'vae')] |
| missing = [path for path in required if not (root / path).is_file()] |
| if missing: |
| raise ValueError(f'Missing model components/configs: {missing}') |
| model_index = json.loads((root / 'model_index.json').read_text(encoding='utf-8')) |
| if model_index.get('_class_name') != 'QwenImage21Pipeline': |
| raise ValueError('Unexpected pipeline class') |
| flags = [] |
| for component in ('transformer', 'text_encoder', 'vae'): |
| folder = root / component |
| config = json.loads((folder / 'config.json').read_text(encoding='utf-8')) |
| files = sorted(folder.glob('*.safetensors')) |
| indexes = list(folder.glob('*.safetensors.index.json')) |
| if not files or len(indexes) > 1 or (len(files) > 1 and not indexes): |
| raise ValueError(f'Missing or ambiguous shards: {component}') |
| index = json.loads(indexes[0].read_text(encoding='utf-8'))['weight_map'] if indexes else None |
| if index is not None and set(index.values()) != {file.name for file in files}: |
| raise ValueError(f'Incomplete shard index: {component}') |
| names = set() |
| for file in files: |
| from safetensors import safe_open |
| with safe_open(str(file), framework='np') as handle: |
| keys = set(handle.keys()) |
| if not keys or names & keys: |
| raise ValueError(f'Empty or duplicate tensors: {file}') |
| if index is not None and any(index.get(key) != file.name for key in keys): |
| raise ValueError(f'Shard tensor mapping mismatch: {file}') |
| names |= keys |
| if index is not None and names != set(index): |
| raise ValueError(f'Missing indexed tensors: {component}') |
| quantized = _quantization(config) |
| if component == 'vae' and quantized: |
| raise ValueError('VAE must remain floating point') |
| if component != 'vae': |
| flags.append(quantized) |
| if flags not in ([False, False], [True, True]): |
| raise ValueError('Transformer and text encoder must use the same precision') |
| return 'int4' if all(flags) else 'bf16' |
|
|
|
|
| def inference_files(root): |
| root = Path(root) |
| result = [root / 'model_index.json'] |
| for component in ('transformer', 'text_encoder', 'vae', 'processor', 'scheduler'): |
| result.extend(path for path in (root / component).rglob('*') |
| if path.is_file() and not any(part.startswith('.') for part in path.relative_to(root).parts)) |
| return sorted(result) |
|
|
|
|
| def model_identity(root, verified_rows=None): |
| root = Path(root) |
| kind = validate_model_structure(root) |
| known = {row['path']: row for row in verified_rows} if verified_rows is not None else None |
| rows = [] |
| for file in inference_files(root): |
| name = file.relative_to(root).as_posix() |
| if known is not None: |
| if name not in known or file.stat().st_size != known[name]['size']: |
| raise ValueError(f'Preverified inventory mismatch: {name}') |
| row = known[name] |
| else: |
| row = {'path': name, 'size': file.stat().st_size, 'sha256': sha256(file)} |
| rows.append(row) |
| encoded = json.dumps(rows, sort_keys=True, separators=(',', ':')).encode() |
| return {'kind': kind, 'base_revision': REVISION, 'files': rows, |
| 'fingerprint': hashlib.sha256(encoded).hexdigest()} |
|
|
|
|
| def validate_roles(baseline, quantized): |
| if baseline['kind'] != 'bf16' or quantized['kind'] != 'int4': |
| raise ValueError('Comparison requires a BF16 baseline and an INT4 candidate') |
| if baseline['fingerprint'] == quantized['fingerprint']: |
| raise ValueError('Baseline and candidate are identical') |
| if baseline['base_revision'] != quantized['base_revision']: |
| raise ValueError('Different upstream revisions') |
|
|
