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