File size: 1,174 Bytes
ffc54ec | 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 | """Shared workstation paths and BF16 helpers for the Image 2.1 Turbo quantization project."""
import os
from pathlib import Path
import torch
from torch import nn
ROOT = Path(os.environ.get('TURBO_ROOT', '/home/user/.local/share/rtx-pro-apps/qwen-image-2.1-turbo'))
BASE = ROOT / 'base'
BASE_REPO = 'Qwen/Qwen-Image-2.1-Turbo'
BASE_REVISION = 'd65dbc9a7e8f6b5479e33dee6030eaab2a906509'
CALIBRATION = ROOT / 'calibration'
def eligible_modules(model):
return {n: m for n, m in model.named_modules()
if isinstance(m, nn.Linear) and n.startswith('transformer_blocks.') and m.bias is None}
def load_bf16_pipeline(base=BASE):
from diffusers import QwenImage21Pipeline
pipe = QwenImage21Pipeline.from_pretrained(base, torch_dtype=torch.bfloat16).to('cuda')
pipe.set_progress_bar_config(disable=True)
return pipe
def sizes(i, kind='calibration'):
"""Same resolution mix as the Image 2.1 releases: every fourth case at 2048, two alternate aspects."""
w = h = 2048 if i % 4 == 0 else 1024
alt = (52, 53) if kind == 'calibration' else (12, 13)
if i == alt[0]: w, h = 1536, 864
if i == alt[1]: w, h = 864, 1536
return w, h
|