live-app-translator/src/livesub/ui/pages/models_page.py

"""모델 — 5단계 품질 티어 선택과 GPU 상태."""

from __future__ import annotations

from PySide6.QtCore import Qt, Signal
from PySide6.QtWidgets import (
    QButtonGroup,
    QComboBox,
    QFrame,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QPushButton,
    QRadioButton,
    QVBoxLayout,
    QWidget,
)

from ...config import AppConfig
from ...models import Tier, detect_gpu, ordered_tiers
from ...models.manager import GpuInfo, tier_availability
from ..theme import SPACING, palette
from ..widgets import Card, PageHeader


class TierCard(QFrame):
    """티어 하나를 라디오 버튼 카드로."""

    def __init__(self, tier: Tier, gpu: GpuInfo, parent: QWidget | None = None):
        super().__init__(parent)
        self.tier = tier
        gpu_vram_gb = gpu.total_vram_mb / 1024 if gpu.available else 0.0
        self.available, self.unavailable_reason = tier_availability(tier, gpu)
        self.setObjectName("Card")
        p = palette("dark")

        layout = QHBoxLayout(self)
        layout.setContentsMargins(16, 14, 16, 14)
        layout.setSpacing(14)

        self.radio = QRadioButton()
        layout.addWidget(self.radio, 0, Qt.AlignmentFlag.AlignTop)

        text = QVBoxLayout()
        text.setSpacing(3)

        head = QHBoxLayout()
        head.setSpacing(8)
        name = QLabel(f"{tier.order}. {tier.name}")
        name.setStyleSheet("font-size: 16px; font-weight: 700;")
        head.addWidget(name)
        if tier.recommended:
            head.addWidget(_badge("추천", p.accent))
        if tier.best_after_finetune:
            head.addWidget(_badge("추가학습 최적", p.success))
        if gpu_vram_gb and gpu_vram_gb < tier.min_vram_gb:
            head.addWidget(_badge(f"VRAM {tier.min_vram_gb:g}GB 필요", p.warning))
        if not self.available:
            head.addWidget(_badge("사용 불가", p.danger))
        head.addStretch(1)
        text.addLayout(head)

        tagline = QLabel(tier.tagline)
        tagline.setObjectName("Subtitle")
        text.addWidget(tagline)

        spec = QLabel(
            f"예상 지연 {tier.approx_latency_s:g}초 · 권장 VRAM {tier.min_vram_gb:g}GB "
            f"· 내려받기 약 {tier.download_gb:g}GB"
        )
        spec.setObjectName("Hint")
        text.addWidget(spec)

        if tier.notes:
            note = QLabel(tier.notes)
            note.setObjectName("Hint")
            note.setWordWrap(True)
            text.addWidget(note)

        if not self.available:
            # 포터블 exe 에는 torch 가 없어 LLM 티어를 못 쓴다. 눌러도 안 되는
            # 이유를 카드에 적어둬야 사용자가 헤매지 않는다.
            blocked = QLabel(self.unavailable_reason)
            blocked.setStyleSheet(f"color: {p.warning};")
            blocked.setWordWrap(True)
            text.addWidget(blocked)
            self.radio.setEnabled(False)
            self.setEnabled(True)

        layout.addLayout(text, 1)

    def set_selected(self, selected: bool) -> None:
        self.setObjectName("CardSelected" if selected else "Card")
        self.style().unpolish(self)
        self.style().polish(self)


class ModelsPage(QWidget):
    tier_changed = Signal(str)
    config_changed = Signal()

    def __init__(self, config: AppConfig, parent: QWidget | None = None):
        super().__init__(parent)
        self.config = config
        self._cards: list[TierCard] = []
        # 티어 카드마다 조회하면 nvidia-smi 를 다섯 번 부른다. 한 번만 본다.
        self._gpu = detect_gpu()
        self._gpu_vram_gb = self._gpu.total_vram_mb / 1024 if self._gpu.available else 0.0

        root = QVBoxLayout(self)
        root.setContentsMargins(0, 0, 0, 0)
        root.setSpacing(SPACING + 4)
        root.addWidget(
            PageHeader(
                "모델",
                "속도 우선(1)부터 품질 우선(5)까지 다섯 단계입니다. "
                "모델은 처음 사용할 때 자동으로 내려받습니다.",
            )
        )

        root.addWidget(self._build_gpu_card())

        self._group = QButtonGroup(self)
        self._group.setExclusive(True)
        for tier in ordered_tiers():
            card = TierCard(tier, self._gpu)
            self._group.addButton(card.radio, tier.order)
            card.radio.toggled.connect(
                lambda checked, t=tier: self._on_tier_selected(t) if checked else None
            )
            self._cards.append(card)
            root.addWidget(card)

        root.addWidget(self._build_lora_card())
        root.addStretch(1)
        self._apply_selection(config.models.tier)

    # --- 구성 -----------------------------------------------------------
    def _build_gpu_card(self) -> Card:
        gpu = self._gpu
        p = palette("dark")

        card = Card("실행 환경")
        row = QWidget()
        layout = QHBoxLayout(row)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(12)

        if gpu.available:
            text = (
                f"GPU: {gpu.name} · VRAM {gpu.total_vram_mb / 1024:.1f}GB "
                f"(여유 {gpu.free_vram_mb / 1024:.1f}GB)"
            )
            color = p.success
        else:
            text = f"GPU를 쓸 수 없습니다 — {gpu.reason} (CPU로도 동작하지만 많이 느립니다)"
            color = p.warning
        info = QLabel(text)
        info.setStyleSheet(f"color: {color}; font-weight: 600;")
        info.setWordWrap(True)
        layout.addWidget(info, 1)

        self.device_combo = QComboBox()
        self.device_combo.addItem("GPU (CUDA)", "cuda")
        self.device_combo.addItem("CPU", "cpu")
        self.device_combo.setCurrentIndex(
            max(0, self.device_combo.findData(self.config.models.device))
        )
        self.device_combo.setEnabled(gpu.available)
        self.device_combo.currentIndexChanged.connect(self._on_device_changed)
        layout.addWidget(self.device_combo)

        card.add(row)
        return card

    def _build_lora_card(self) -> Card:
        card = Card(
            "추가학습 어댑터 (선택)",
            "게임·방송 용어로 따로 학습시킨 LoRA 어댑터 폴더를 지정하면 번역 모델에 얹습니다. "
            "scripts/finetune_mt.py 로 만들 수 있습니다.",
        )
        row = QWidget()
        layout = QHBoxLayout(row)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(10)

        self.lora_edit = QLineEdit(self.config.models.lora_adapter_path)
        self.lora_edit.setPlaceholderText("비워두면 사용하지 않습니다")
        self.lora_edit.editingFinished.connect(self._on_lora_changed)
        layout.addWidget(self.lora_edit, 1)

        browse = QPushButton("폴더 선택")
        browse.clicked.connect(self._browse_lora)
        layout.addWidget(browse)

        card.add(row)
        return card

    # --- 동작 -----------------------------------------------------------
    def _apply_selection(self, tier_key: str) -> None:
        for card in self._cards:
            selected = card.tier.key == tier_key
            card.radio.setChecked(selected)
            card.set_selected(selected)

    def _on_tier_selected(self, tier: Tier) -> None:
        if self.config.models.tier == tier.key:
            return
        self.config.models.tier = tier.key
        for card in self._cards:
            card.set_selected(card.tier.key == tier.key)
        self.tier_changed.emit(tier.key)
        self.config_changed.emit()

    def _on_device_changed(self) -> None:
        self.config.models.device = self.device_combo.currentData()
        self.config_changed.emit()

    def _on_lora_changed(self) -> None:
        self.config.models.lora_adapter_path = self.lora_edit.text().strip()
        self.config_changed.emit()

    def _browse_lora(self) -> None:
        from PySide6.QtWidgets import QFileDialog

        path = QFileDialog.getExistingDirectory(self, "LoRA 어댑터 폴더 선택")
        if path:
            self.lora_edit.setText(path)
            self._on_lora_changed()


def _badge(text: str, color: str) -> QLabel:
    label = QLabel(text)
    label.setStyleSheet(
        f"color: {color}; border: 1px solid {color}; border-radius: 9px;"
        f"padding: 1px 8px; font-size: 11px; font-weight: 700;"
    )
    return label