File size: 11,679 Bytes
e0265b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
from __future__ import annotations

from pathlib import Path

from PySide6.QtCore import Qt, QUrl
from PySide6.QtGui import QDesktopServices
from PySide6.QtWidgets import (
    QAbstractItemView, QComboBox, QFrame, QGridLayout, QHBoxLayout, QLabel,
    QLineEdit, QListWidget, QListWidgetItem, QMessageBox, QPushButton,
    QSpinBox, QVBoxLayout, QWidget,
)

from adam.assets import AssetRegistry
from adam.config import ConfigManager
from adam.generations import build_generation_plan, generation_tools
from adam.job_manager import JobManager
from adam.models import Job, JobStatus
from adam.registry import ToolRegistry
from adam.showcase import build_showcase_plan


def _card() -> QFrame:
    card = QFrame(); card.setProperty("card", True); return card


def _title(text: str) -> QLabel:
    label = QLabel(text); label.setObjectName("CardTitle"); return label


class ShowcasePage(QWidget):
    """Create a finished, self-generated DDPM/Flow showcase video."""

    def __init__(self, root: Path, registry: ToolRegistry, jobs: JobManager, assets: AssetRegistry, config: ConfigManager) -> None:
        super().__init__()
        self.root = root.resolve(); self.registry = registry; self.jobs = jobs; self.assets = assets; self.config = config
        self.latest_video: Path | None = None
        layout = QVBoxLayout(self); layout.setContentsMargins(22, 18, 22, 18); layout.setSpacing(12)
        heading = QLabel("Showcase Video"); heading.setObjectName("PageTitle")
        subtitle = QLabel("Generate 12–24 images per DDPM or Flow Matching model, then export the animated request interface as a finished MP4.")
        subtitle.setProperty("muted", True); subtitle.setWordWrap(True)
        layout.addWidget(heading); layout.addWidget(subtitle)
        body = QHBoxLayout(); body.setSpacing(12)
        body.addWidget(self._build_models(), 1); body.addWidget(self._build_settings(), 1)
        layout.addLayout(body, 1)
        self.jobs.job_updated.connect(self._job_updated)
        self.refresh(); self._restore()

    def _build_models(self) -> QFrame:
        card = _card(); layout = QVBoxLayout(card); layout.setContentsMargins(16, 15, 16, 16)
        top = QHBoxLayout(); top.addWidget(_title("REQUEST / MODEL ORDER")); top.addStretch()
        refresh = QPushButton("Refresh"); refresh.clicked.connect(self.refresh); top.addWidget(refresh); layout.addLayout(top)
        note = QLabel("Select one or more completed models. Drag entries to set their order in the video. LoRAs are intentionally excluded.")
        note.setProperty("muted", True); note.setWordWrap(True); layout.addWidget(note)
        self.models = QListWidget(); self.models.setSelectionMode(QAbstractItemView.MultiSelection); self.models.setDragDropMode(QAbstractItemView.InternalMove); self.models.setDefaultDropAction(Qt.MoveAction); self.models.setSpacing(5)
        layout.addWidget(self.models, 1)
        buttons = QHBoxLayout(); select_all = QPushButton("Select all"); clear = QPushButton("Clear")
        select_all.clicked.connect(self.models.selectAll); clear.clicked.connect(self.models.clearSelection)
        buttons.addWidget(select_all); buttons.addWidget(clear); buttons.addStretch(); layout.addLayout(buttons)
        return card

    def _build_settings(self) -> QFrame:
        card = _card(); layout = QVBoxLayout(card); layout.setContentsMargins(16, 15, 16, 16); layout.setSpacing(10)
        layout.addWidget(_title("SHOWCASE SETTINGS"))
        grid = QGridLayout(); grid.setVerticalSpacing(8)
        self.video_title = QLineEdit("ADAM Generation Showcase")
        self.images = QSpinBox(); self.images.setRange(12, 24); self.images.setValue(18); self.images.setSuffix(" per model")
        self.duration = QComboBox(); self.duration.addItems(["3 seconds", "4 seconds", "5 seconds"]); self.duration.setCurrentIndex(1)
        self.steps = QSpinBox(); self.steps.setRange(5, 200); self.steps.setValue(30)
        self.ddpm_sampler = QComboBox(); self.ddpm_sampler.addItems(["DDIM", "DDPM"])
        self.flow_sampler = QComboBox(); self.flow_sampler.addItems(["Heun", "Euler"])
        self.aspect = QComboBox(); self.aspect.addItems(["16:9 (Widescreen)", "1:1 (Square)"])
        self.seed = QSpinBox(); self.seed.setRange(1, 2_146_000_000); self.seed.setValue(123456)
        self.resolution = QComboBox(); self.resolution.addItems(["1080p", "720p"])
        fields = [
            ("Video title", self.video_title), ("Images", self.images), ("Each image lasts", self.duration),
            ("Generation steps", self.steps), ("DDPM sampler", self.ddpm_sampler),
            ("Flow method", self.flow_sampler), ("Image aspect", self.aspect),
            ("Starting seed", self.seed), ("Video resolution", self.resolution),
        ]
        for row, (label, widget) in enumerate(fields): grid.addWidget(QLabel(label), row, 0); grid.addWidget(widget, row, 1)
        layout.addLayout(grid)
        self.estimate = QLabel(); self.estimate.setProperty("muted", True); self.estimate.setWordWrap(True); layout.addWidget(self.estimate)
        self.create_button = QPushButton("Generate images and create MP4  →"); self.create_button.setProperty("primary", True); self.create_button.clicked.connect(self._create)
        layout.addWidget(self.create_button)
        self.open_button = QPushButton("Open latest showcase video"); self.open_button.setEnabled(False); self.open_button.clicked.connect(self._open_latest); layout.addWidget(self.open_button)
        self.status = QLabel("Ready"); self.status.setProperty("muted", True); self.status.setWordWrap(True); layout.addWidget(self.status)
        layout.addStretch()
        self.models.itemSelectionChanged.connect(self._update_estimate)
        self.images.valueChanged.connect(self._update_estimate); self.duration.currentIndexChanged.connect(self._update_estimate)
        return card

    @staticmethod
    def _ready(path: Path, trainer: str) -> bool:
        if trainer == "ddpm": return path.is_dir() and (path / "model_index.json").is_file()
        if trainer == "flow": return path.is_dir() and (path / "flow_model_info.json").is_file() and (path / "unet" / "config.json").is_file()
        return False

    def refresh(self) -> None:
        selected = {str(item.data(Qt.UserRole).get("path")) for item in self.models.selectedItems()} if hasattr(self, "models") else set()
        self.assets.discover(self.config); self.models.clear()
        tools = generation_tools(self.registry)
        for asset in self.assets.assets:
            if asset.kind != "model" or asset.trainer not in {"ddpm", "flow"} or not self._ready(Path(asset.path), asset.trainer): continue
            tool = next((value for value in tools if asset.trainer in value.model_trainers), None)
            if not tool: continue
            label = "DDPM" if asset.trainer == "ddpm" else "Flow Matching"
            item = QListWidgetItem(f"{asset.name}{label}")
            item.setData(Qt.UserRole, {"name": asset.name, "path": asset.path, "trainer": asset.trainer, "trainer_label": label, "tool_id": tool.id})
            self.models.addItem(item); item.setSelected(asset.path in selected)
        self.create_button.setEnabled(self.models.count() > 0); self._update_estimate()

    def _duration_seconds(self) -> int:
        return int(self.duration.currentText().split()[0])

    def _update_estimate(self) -> None:
        count = len(self.models.selectedItems()) if hasattr(self, "models") else 0
        images = count * self.images.value() if hasattr(self, "images") else 0
        seconds = images * self._duration_seconds() if images else 0
        self.estimate.setText(f"{images} images total · finished video about {seconds // 60}:{seconds % 60:02d}" if count else "Select at least one model to calculate the video length.")

    def _selected(self) -> list[dict]:
        return [dict(self.models.item(index).data(Qt.UserRole)) for index in range(self.models.count()) if self.models.item(index).isSelected()]

    def _create(self) -> None:
        selected = self._selected()
        if not selected:
            QMessageBox.information(self, "Choose models", "Select at least one completed DDPM or Flow Matching model."); return
        count = self.images.value(); plans = []; settings = []
        for index, entry in enumerate(selected):
            tool = self.registry.get(entry["tool_id"]); trainer = entry["trainer"]
            sampler = self.ddpm_sampler.currentText() if trainer == "ddpm" else self.flow_sampler.currentText()
            plans.append(build_generation_plan(
                tool, model_name=entry["name"], model_path=entry["path"], prompt="",
                image_count=count, steps=self.steps.value(), seed=self.seed.value() + index * count,
                sampler=sampler, aspect_ratio=self.aspect.currentText(),
            ))
            settings.append({**entry, "steps": self.steps.value(), "sampler": sampler, "aspect_ratio": self.aspect.currentText()})
        plan = build_showcase_plan(
            plans, title=self.video_title.text(), display_seconds=self._duration_seconds(),
            resolution=self.resolution.currentText(), model_settings=settings,
        )
        job = self.jobs.submit(plan); self._save(); self.status.setText(f"Showcase job {job.id} queued. ADAM will generate every image before rendering the MP4.")

    def _save(self) -> None:
        self.config.update({"showcase_settings": {"title": self.video_title.text(), "images": self.images.value(), "duration": self._duration_seconds(), "steps": self.steps.value(), "ddpm_sampler": self.ddpm_sampler.currentText(), "flow_sampler": self.flow_sampler.currentText(), "aspect": self.aspect.currentText(), "seed": self.seed.value(), "resolution": self.resolution.currentText()}})

    def _restore(self) -> None:
        saved = self.config.get("showcase_settings", {})
        if not isinstance(saved, dict): return
        self.video_title.setText(str(saved.get("title", self.video_title.text())))
        for spin, key in ((self.images, "images"), (self.steps, "steps"), (self.seed, "seed")):
            try: spin.setValue(max(spin.minimum(), min(int(saved.get(key, spin.value())), spin.maximum())))
            except (TypeError, ValueError): pass
        duration = f"{saved.get('duration', 4)} seconds"
        for combo, value in ((self.duration, duration), (self.ddpm_sampler, saved.get("ddpm_sampler")), (self.flow_sampler, saved.get("flow_sampler")), (self.aspect, saved.get("aspect")), (self.resolution, saved.get("resolution"))):
            if value and combo.findText(str(value)) >= 0: combo.setCurrentText(str(value))
        self._update_estimate()

    def _job_updated(self, job: Job) -> None:
        if job.plan.project_name != "Showcase Video": return
        if job.status == JobStatus.FINISHED:
            videos = sorted((self.root / "data" / "showcase_videos").glob(f"*{job.id}.mp4"))
            self.latest_video = videos[-1] if videos else None; self.open_button.setEnabled(self.latest_video is not None)
            self.status.setText(f"Showcase {job.id} finished successfully." if self.latest_video else f"Showcase {job.id} finished, but its MP4 could not be located.")
        elif job.status == JobStatus.FAILED: self.status.setText(f"Showcase {job.id} failed: {job.error or 'Unknown error'}")
        elif job.status == JobStatus.CANCELLED: self.status.setText(f"Showcase {job.id} was cancelled.")
        else: self.status.setText(f"Showcase {job.id}: {job.status.value} · {job.progress}%")

    def _open_latest(self) -> None:
        if self.latest_video and self.latest_video.is_file(): QDesktopServices.openUrl(QUrl.fromLocalFile(str(self.latest_video)))