from __future__ import annotations import json from pathlib import Path from PySide6.QtCore import QSize, Qt, QTimer, QUrl from PySide6.QtGui import QDesktopServices, QIcon, QImageReader, QPixmap from PySide6.QtWidgets import ( QAbstractItemView, QComboBox, QCheckBox, QDoubleSpinBox, QDialog, QDialogButtonBox, QFileDialog, QFrame, QGridLayout, QHBoxLayout, QLabel, QLineEdit, QListWidget, QListWidgetItem, QMessageBox, QPlainTextEdit, QPushButton, QScrollArea, QSplitter, QSpinBox, QVBoxLayout, QWidget, ) from adam.assets import Asset, AssetRegistry from adam.config import ConfigManager from adam.generations import ( GenerationRecord, build_generation_plan, generation_tools, load_generation_history, combine_generation_plans, ) from adam.job_manager import JobManager from adam.models import Job, JobStatus from adam.registry import ToolRegistry, ToolSpec def _card() -> QFrame: frame = QFrame() frame.setProperty("card", True) return frame def _title(text: str) -> QLabel: label = QLabel(text) label.setObjectName("CardTitle") return label def _header(title: str, subtitle: str) -> QWidget: widget = QWidget() layout = QVBoxLayout(widget) layout.setContentsMargins(0, 0, 0, 10) heading = QLabel(title) heading.setObjectName("PageTitle") detail = QLabel(subtitle) detail.setProperty("muted", True) detail.setWordWrap(True) layout.addWidget(heading) layout.addWidget(detail) return widget def _thumbnail(path: Path, width: int, height: int) -> QPixmap: reader = QImageReader(str(path)) reader.setAutoTransform(True) source_size = reader.size() if source_size.isValid(): source_size.scale(QSize(width, height), Qt.KeepAspectRatio) reader.setScaledSize(source_size) image = reader.read() return QPixmap.fromImage(image) if not image.isNull() else QPixmap() class GenerationCycleDialog(QDialog): """Choose several completed models and shared playback settings.""" def __init__(self, entries: list[dict], parent: QWidget | None = None) -> None: super().__init__(parent) self.entries = entries self.setWindowTitle("Generation Cycle") self.setMinimumSize(600, 560) layout = QVBoxLayout(self) layout.addWidget(_header( "Generation Cycle", "Select completed models. ADAM will generate each batch sequentially, then play the results as a labeled slideshow.", )) self.models = QListWidget() self.models.setSelectionMode(QAbstractItemView.MultiSelection) for index, entry in enumerate(entries): item = QListWidgetItem(f"{entry['name']} — {entry['trainer_label']}") item.setData(Qt.UserRole, index) self.models.addItem(item) selection_buttons = QHBoxLayout() select_all = QPushButton("Select all") clear = QPushButton("Clear") select_all.clicked.connect(self.models.selectAll) clear.clicked.connect(self.models.clearSelection) selection_buttons.addWidget(select_all); selection_buttons.addWidget(clear); selection_buttons.addStretch() layout.addLayout(selection_buttons) layout.addWidget(self.models, 1) form = QGridLayout() self.image_count = QSpinBox(); self.image_count.setRange(1, 48); self.image_count.setValue(4) self.display_seconds = QSpinBox(); self.display_seconds.setRange(1, 300); self.display_seconds.setValue(5); self.display_seconds.setSuffix(" seconds") self.seed = QSpinBox(); self.seed.setRange(0, 2_147_000_000); self.seed.setValue(123456) self.prompt = QLineEdit(); self.prompt.setPlaceholderText("Optional shared prompt or creative note") self.show_labels = QCheckBox("Show model and trainer in the bottom-left corner"); self.show_labels.setChecked(True) self.loop = QCheckBox("Loop the slideshow until closed") self.fullscreen = QCheckBox("Start playback in fullscreen") for row, (label, widget) in enumerate(( ("Images per model", self.image_count), ("Display each image", self.display_seconds), ("Starting seed", self.seed), ("Shared prompt / note", self.prompt), )): form.addWidget(QLabel(label), row, 0); form.addWidget(widget, row, 1) layout.addLayout(form) layout.addWidget(self.show_labels); layout.addWidget(self.loop); layout.addWidget(self.fullscreen) self.validation = QLabel(); self.validation.setProperty("muted", True); self.validation.setWordWrap(True) layout.addWidget(self.validation) buttons = QDialogButtonBox(QDialogButtonBox.Cancel | QDialogButtonBox.Ok) buttons.button(QDialogButtonBox.Ok).setText("Queue generation cycle") buttons.accepted.connect(self._accept); buttons.rejected.connect(self.reject) layout.addWidget(buttons) def _accept(self) -> None: if not self.models.selectedItems(): self.validation.setText("Select at least one completed model.") return self.accept() def selected_entries(self) -> list[dict]: selected = {int(item.data(Qt.UserRole)) for item in self.models.selectedItems()} return [entry for index, entry in enumerate(self.entries) if index in selected] class CycleSlideshowDialog(QDialog): """Simple local slideshow for a completed multi-model generation cycle.""" def __init__( self, slides: list[tuple[Path, str]], *, display_seconds: int, show_labels: bool, loop: bool, parent: QWidget | None = None, ) -> None: super().__init__(parent) self.slides = slides self.index = 0 self.show_labels = show_labels self.loop = loop self.setWindowTitle("Generation Cycle Playback") self.setMinimumSize(800, 560) self.setStyleSheet("background: #02070b;") layout = QVBoxLayout(self) layout.setContentsMargins(0, 0, 0, 0) self.image = QLabel() self.image.setAlignment(Qt.AlignCenter) self.image.setMinimumSize(500, 350) self.overlay = QLabel(self.image) self.overlay.setStyleSheet( "background: rgba(0, 0, 0, 180); color: white; padding: 8px 12px; " "border-radius: 5px; font-weight: 600;" ) layout.addWidget(self.image, 1) controls = QHBoxLayout() self.previous = QPushButton("Previous") self.caption = QLabel(); self.caption.setAlignment(Qt.AlignCenter) self.next = QPushButton("Next") close = QPushButton("Close") controls.addWidget(self.previous); controls.addWidget(self.caption, 1); controls.addWidget(self.next); controls.addWidget(close) layout.addLayout(controls) self.previous.clicked.connect(self._previous); self.next.clicked.connect(self._next); close.clicked.connect(self.accept) self.timer = QTimer(self) self.timer.setInterval(max(1, int(display_seconds)) * 1000) self.timer.timeout.connect(self._next) self._show_slide() self.timer.start() def resizeEvent(self, event) -> None: super().resizeEvent(event) self._show_slide() def _show_slide(self) -> None: if not self.slides: return path, label = self.slides[self.index] pixmap = QPixmap(str(path)) if not pixmap.isNull(): self.image.setPixmap(pixmap.scaled(self.image.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)) else: self.image.setText(f"Could not display {path.name}") self.caption.setText(f"Image {self.index + 1} of {len(self.slides)}") self.overlay.setText(label) self.overlay.adjustSize() self.overlay.move(18, max(18, self.image.height() - self.overlay.height() - 18)) self.overlay.setVisible(self.show_labels) def _previous(self) -> None: self.index = (self.index - 1) % len(self.slides) self._show_slide() def _next(self) -> None: if self.index + 1 >= len(self.slides) and not self.loop: self.timer.stop() return self.index = (self.index + 1) % len(self.slides) self._show_slide() class GenerationsPage(QWidget): """Generate with registered providers and revisit reproducible output batches.""" 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.records: list[GenerationRecord] = [] self.hidden_history_images: set[str] = set() self._cycle_jobs: dict[str, dict] = {} layout = QVBoxLayout(self) layout.setContentsMargins(22, 18, 22, 18) layout.setSpacing(12) header_row = QHBoxLayout() header_row.addWidget( _header( "Generations", "Create images with compatible registered tools, then revisit every seed and setting.", ), 1 ) self.cycle_button = QPushButton("Generation Cycle…") self.cycle_button.setProperty("primary", True) header_row.addWidget(self.cycle_button) layout.addLayout(header_row) body = QSplitter(Qt.Horizontal) body.setChildrenCollapsible(False) body.addWidget(self._build_controls()) body.addWidget(self._build_history()) body.setStretchFactor(0, 0) body.setStretchFactor(1, 1) body.setSizes([430, 1000]) layout.addWidget(body, 1) self.provider.currentIndexChanged.connect(self._provider_changed) self.model.currentIndexChanged.connect(self._model_changed) self.preset.currentIndexChanged.connect(self._apply_preset) self.aspect.currentTextChanged.connect(self._lora_aspect_changed) self.generate_button.clicked.connect(self._generate) self.refresh_button.clicked.connect(self.refresh) self.gallery.currentItemChanged.connect(self._selection_changed) self.gallery.itemDoubleClicked.connect(lambda _item: self._open_image()) self.open_image_button.clicked.connect(self._open_image) self.open_folder_button.clicked.connect(self._open_folder) self.reuse_button.clicked.connect(self._reuse_settings) self.clear_history_button.clicked.connect(self._clear_displayed_history) self.random_seed_button.clicked.connect(lambda: self.seed.clear()) self.jobs.job_updated.connect(self._job_updated) self.cycle_button.clicked.connect(self._open_generation_cycle) self.refresh() self._restore_generation_settings() self._connect_generation_setting_persistence() def _build_controls(self) -> QFrame: card = _card() card.setMinimumWidth(340) outer_layout = QVBoxLayout(card) outer_layout.setContentsMargins(0, 0, 0, 0) scroll = QScrollArea() scroll.setWidgetResizable(True) scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) scroll.setFrameShape(QFrame.NoFrame) content = QWidget() scroll.setWidget(content) outer_layout.addWidget(scroll) layout = QVBoxLayout(content) layout.setContentsMargins(16, 15, 16, 16) layout.setSpacing(9) layout.addWidget(_title("CREATE A BATCH")) layout.addWidget(QLabel("Registered generator")) self.provider = QComboBox() layout.addWidget(self.provider) layout.addWidget(QLabel("Completed model")) model_row = QHBoxLayout() self.model = QComboBox() self.refresh_button = QPushButton("Refresh") self.model_browse_button = QPushButton("Browse LoRA") self.model_browse_button.setVisible(False) self.model_browse_button.clicked.connect(self._browse_lora) model_row.addWidget(self.model, 1) model_row.addWidget(self.model_browse_button) model_row.addWidget(self.refresh_button) layout.addLayout(model_row) self.model_hint = QLabel("Select a registered generator and model.") self.model_hint.setProperty("muted", True) self.model_hint.setWordWrap(True) layout.addWidget(self.model_hint) self.prompt_label = QLabel("Label / creative note") layout.addWidget(self.prompt_label) self.prompt = QPlainTextEdit() self.prompt.setMaximumHeight(78) self.prompt.setPlaceholderText( "Optional note describing what you are exploring. Unconditional generators save it with the batch but do not use it as a text prompt." ) layout.addWidget(self.prompt) self.lora_options = QFrame() lora_layout = QGridLayout(self.lora_options) lora_layout.setContentsMargins(0, 0, 0, 0) lora_layout.setHorizontalSpacing(8) lora_layout.setVerticalSpacing(6) lora_layout.addWidget(QLabel("Base model"), 0, 0) self.base_model = QComboBox() self.base_model.setEditable(True) self.base_model.setInsertPolicy(QComboBox.NoInsert) lora_layout.addWidget(self.base_model, 0, 1) lora_layout.addWidget(QLabel("Negative prompt"), 1, 0) self.negative_prompt = QPlainTextEdit() self.negative_prompt.setMaximumHeight(62) lora_layout.addWidget(self.negative_prompt, 1, 1) lora_layout.addWidget(QLabel("Width"), 2, 0) self.width = QSpinBox() self.width.setRange(256, 2048) self.width.setSingleStep(8) lora_layout.addWidget(self.width, 2, 1) lora_layout.addWidget(QLabel("Height"), 3, 0) self.height = QSpinBox() self.height.setRange(256, 2048) self.height.setSingleStep(8) lora_layout.addWidget(self.height, 3, 1) lora_layout.addWidget(QLabel("CFG scale"), 4, 0) self.cfg_scale = QDoubleSpinBox() self.cfg_scale.setRange(0.1, 30.0) self.cfg_scale.setSingleStep(0.5) self.cfg_scale.setDecimals(2) lora_layout.addWidget(self.cfg_scale, 4, 1) lora_layout.addWidget(QLabel("LoRA strength"), 5, 0) self.lora_strength = QDoubleSpinBox() self.lora_strength.setRange(0.05, 3.0) self.lora_strength.setSingleStep(0.05) self.lora_strength.setDecimals(2) lora_layout.addWidget(self.lora_strength, 5, 1) lora_layout.addWidget(QLabel("Reference image"), 6, 0) reference_row = QHBoxLayout() self.reference_image = QLineEdit() self.reference_browse = QPushButton("Browse") self.reference_browse.clicked.connect(self._browse_reference) reference_row.addWidget(self.reference_image, 1) reference_row.addWidget(self.reference_browse) lora_layout.addLayout(reference_row, 6, 1) lora_layout.addWidget(QLabel("Denoise strength"), 7, 0) self.denoise_strength = QDoubleSpinBox() self.denoise_strength.setRange(0.0, 1.0) self.denoise_strength.setSingleStep(0.05) self.denoise_strength.setDecimals(2) lora_layout.addWidget(self.denoise_strength, 7, 1) self.prompt_weighting = QCheckBox("Use CivitAI/A1111 prompt weights") lora_layout.addWidget(self.prompt_weighting, 8, 0, 1, 2) self.lora_options.setVisible(False) layout.addWidget(self.lora_options) self.ddpm_reference_options = QFrame() ddpm_reference_layout = QGridLayout(self.ddpm_reference_options) ddpm_reference_layout.setContentsMargins(0, 0, 0, 0) ddpm_reference_layout.setHorizontalSpacing(8) ddpm_reference_layout.setVerticalSpacing(6) ddpm_reference_layout.addWidget(QLabel("Reference image"), 0, 0) ddpm_reference_row = QHBoxLayout() self.ddpm_reference_image = QLineEdit() self.ddpm_reference_image.setPlaceholderText("Optional image to reimagine") self.ddpm_reference_browse = QPushButton("Browse") self.ddpm_reference_browse.clicked.connect(self._browse_ddpm_reference) ddpm_reference_row.addWidget(self.ddpm_reference_image, 1) ddpm_reference_row.addWidget(self.ddpm_reference_browse) ddpm_reference_layout.addLayout(ddpm_reference_row, 0, 1) ddpm_reference_layout.addWidget(QLabel("Reference strength"), 1, 0) self.ddpm_reference_strength = QSpinBox() self.ddpm_reference_strength.setRange(0, 100) self.ddpm_reference_strength.setSuffix("%") self.ddpm_reference_strength.setValue(65) self.ddpm_reference_strength.setToolTip( "0% mostly ignores the image; 100% preserves it as closely as possible." ) ddpm_reference_layout.addWidget(self.ddpm_reference_strength, 1, 1) self.ddpm_custom_size = QCheckBox("Use custom dimensions") self.ddpm_custom_size.setToolTip( "Use exact dimensions instead of the selected aspect ratio. Both values must match the model's required size multiple." ) self.ddpm_custom_size.toggled.connect(self._ddpm_custom_size_changed) ddpm_reference_layout.addWidget(self.ddpm_custom_size, 2, 0, 1, 2) ddpm_reference_layout.addWidget(QLabel("Width"), 3, 0) self.ddpm_width = QSpinBox() self.ddpm_width.setRange(64, 2048) self.ddpm_width.setSingleStep(8) self.ddpm_width.setValue(256) self.ddpm_width.setSuffix(" px") ddpm_reference_layout.addWidget(self.ddpm_width, 3, 1) ddpm_reference_layout.addWidget(QLabel("Height"), 4, 0) self.ddpm_height = QSpinBox() self.ddpm_height.setRange(64, 2048) self.ddpm_height.setSingleStep(8) self.ddpm_height.setValue(256) self.ddpm_height.setSuffix(" px") ddpm_reference_layout.addWidget(self.ddpm_height, 4, 1) self._ddpm_custom_size_changed(False) self.ddpm_reference_options.setVisible(False) layout.addWidget(self.ddpm_reference_options) grid = QGridLayout() grid.setHorizontalSpacing(8) grid.setVerticalSpacing(6) self.preset = QComboBox() self.preset.addItems(["Balanced", "Fast preview", "Detailed", "Classic"]) self.images = QSpinBox() self.images.setRange(1, 48) self.images.setValue(1) self.steps = QSpinBox() self.steps.setRange(5, 500) self.steps.setValue(50) self.preview_interval = QSpinBox() self.preview_interval.setRange(0, 500) self.preview_interval.setValue(0) self.preview_interval.setSpecialValueText("Off") self.preview_interval.setSuffix(" steps") self.preview_interval.setToolTip( "Decode and show a live preview after this many denoising steps. Smaller intervals can make generation slower." ) self.sampler = QComboBox() self.sampler.addItems(["DDIM", "DDPM"]) self.aspect = QComboBox() self.aspect.addItems( [ "1:1 (Square)", "16:9 (Widescreen)", "9:16 (Portrait)", "4:3 (Classic)", "3:4 (Portrait Classic)", "3:2 (Photo)", "2:3 (Portrait Photo)", ] ) fields = [ ("Preset", self.preset), ("Images", self.images), ("Steps", self.steps), ("Steps / preview", self.preview_interval), ("Sampler", self.sampler), ("Aspect ratio", self.aspect), ] for row, (label, widget) in enumerate(fields): grid.addWidget(QLabel(label), row, 0) grid.addWidget(widget, row, 1) layout.addLayout(grid) layout.addWidget(QLabel("Seed")) seed_row = QHBoxLayout() self.seed = QLineEdit() self.seed.setPlaceholderText("Random") self.random_seed_button = QPushButton("Randomize") seed_row.addWidget(self.seed, 1) seed_row.addWidget(self.random_seed_button) layout.addLayout(seed_row) self.generate_button = QPushButton("Generate images →") self.generate_button.setProperty("primary", True) layout.addWidget(self.generate_button) self.status = QLabel("Ready") self.status.setProperty("muted", True) self.status.setWordWrap(True) layout.addWidget(self.status) layout.addStretch() return card def _build_history(self) -> QFrame: card = _card() layout = QVBoxLayout(card) layout.setContentsMargins(16, 15, 16, 16) layout.setSpacing(9) top = QHBoxLayout() top.addWidget(_title("GENERATION HISTORY")) top.addStretch() self.history_summary = QLabel("No generations yet") self.history_summary.setProperty("muted", True) top.addWidget(self.history_summary) self.clear_history_button = QPushButton("Clear displayed") self.clear_history_button.setToolTip( "Hide the displayed thumbnails without deleting generation records or image files." ) self.clear_history_button.setEnabled(False) top.addWidget(self.clear_history_button) layout.addLayout(top) self.gallery = QListWidget() self.gallery.setViewMode(QListWidget.IconMode) self.gallery.setIconSize(QSize(170, 128)) self.gallery.setGridSize(QSize(205, 188)) self.gallery.setResizeMode(QListWidget.Adjust) self.gallery.setSelectionMode(QAbstractItemView.SingleSelection) layout.addWidget(self.gallery, 1) self.detail = QLabel("Select an image to see its reproducibility settings.") self.detail.setProperty("muted", True) self.detail.setWordWrap(True) layout.addWidget(self.detail) actions = QHBoxLayout() self.open_image_button = QPushButton("Open image") self.open_folder_button = QPushButton("Open folder") self.reuse_button = QPushButton("Reuse settings") for button in ( self.open_image_button, self.open_folder_button, self.reuse_button, ): button.setEnabled(False) actions.addWidget(button) actions.addStretch() layout.addLayout(actions) return card def refresh(self) -> None: selected_provider = self.provider.currentData() selected_model = self.model.currentData() self.assets.discover(self.config) self.provider.blockSignals(True) self.provider.clear() for tool in generation_tools(self.registry): self.provider.addItem(tool.name, tool.id) if selected_provider: index = self.provider.findData(selected_provider) if index >= 0: self.provider.setCurrentIndex(index) self.provider.blockSignals(False) self._provider_changed() if selected_model: index = self.model.findData(selected_model) if index >= 0: self.model.setCurrentIndex(index) self._load_history() def _current_tool(self) -> ToolSpec | None: tool_id = self.provider.currentData() if not tool_id: return None try: return self.registry.get(str(tool_id)) except Exception: return None def _provider_changed(self) -> None: tool = self._current_tool() current_path = self.model.currentData() self.model.blockSignals(True) self.model.clear() if tool: candidates = [ asset for asset in self.assets.assets if asset.kind == "model" and (not tool.model_trainers or asset.trainer in tool.model_trainers) and self._model_is_ready(asset) ] for asset in candidates: self.model.addItem(asset.name, asset.path) if current_path: index = self.model.findData(current_path) if index >= 0: self.model.setCurrentIndex(index) options = tool.generation_options samplers = [str(item) for item in options.get("samplers", [])] aspects = [str(item) for item in options.get("aspect_ratios", [])] self.sampler.clear() self.sampler.addItems(samplers or ["DDIM", "DDPM"]) self.aspect.clear() self.aspect.addItems(aspects or ["1:1 (Square)"]) step_min = int(options.get("step_min", 5) or 5) step_max = int(options.get("step_max", 500) or 500) self.steps.setRange(step_min, step_max) self.steps.setValue( max(step_min, min(int(options.get("step_default", 50) or 50), step_max)) ) self.preview_interval.setRange(0, step_max) is_prompted = "text_prompt" in tool.capabilities self.model_browse_button.setVisible(tool.id == "lora_generator") self.prompt_label.setText("Positive prompt" if is_prompted else "Label / creative note") self.prompt.setPlaceholderText( "Describe the image to create with the selected LoRA." if is_prompted else "Optional note describing what you are exploring. Unconditional generators save it with the batch but do not use it as a text prompt." ) self.lora_options.setVisible(tool.id == "lora_generator") self.ddpm_reference_options.setVisible(tool.id == "ddpm_generator") self.aspect.setEnabled( tool.id != "ddpm_generator" or not self.ddpm_custom_size.isChecked() ) if tool.id == "lora_generator": self._load_lora_options() self.model.blockSignals(False) self.generate_button.setEnabled(bool(tool and self.model.count())) self._apply_preset() self._model_changed() def _browse_lora(self) -> None: """Choose a LoRA outside the connected trainer's output folder.""" path, _ = QFileDialog.getOpenFileName( self, "Choose LoRA file", "", "LoRA files (*.safetensors)", ) if not path: return selected = Path(path).resolve() if "_comfy" in selected.stem.casefold(): QMessageBox.warning( self, "Choose the original LoRA", "The *_comfy.safetensors export is for ComfyUI. Choose the original LoRA file instead.", ) return index = self.model.findData(str(selected)) if index < 0: self.model.addItem(selected.stem, str(selected)) index = self.model.count() - 1 self.model.setCurrentIndex(index) def _lora_settings(self) -> dict: root = Path(str(self.config.get("tool_folders", {}).get("lora_trainer", ""))) try: value = json.loads((root / "config" / "app_settings.json").read_text(encoding="utf-8")) return value if isinstance(value, dict) else {} except (OSError, ValueError, TypeError, json.JSONDecodeError): return {} def _load_lora_options(self) -> None: saved = self._lora_settings() base_models = [str(saved.get(key, "")).strip() for key in ("generate_model", "last_model")] current = self.base_model.currentText().strip() self.base_model.clear() preferred_asset = next( ( asset for asset in self.assets.assets if asset.kind == "base_model" and "waiillustrious" in "".join( character for character in asset.name.casefold() if character.isalnum() ) ), None, ) if preferred_asset is not None: self.base_model.addItem(preferred_asset.name, preferred_asset.path) for path in dict.fromkeys(item for item in base_models if item): self.base_model.addItem(Path(path).name, path) for asset in self.assets.assets: if asset.kind == "base_model" and self.base_model.findData(asset.path) < 0: self.base_model.addItem(asset.name, asset.path) if current: self.base_model.setCurrentText(current) elif preferred_asset is not None: self.base_model.setCurrentIndex(0) elif base_models and base_models[0]: self.base_model.setCurrentText(base_models[0]) self.negative_prompt.setPlainText(str(saved.get("generate_negative_prompt", ""))) self.cfg_scale.setValue(float(saved.get("generate_cfg", 7.0) or 7.0)) self.lora_strength.setValue(float(saved.get("generate_lora_strength", 1.0) or 1.0)) self.denoise_strength.setValue(float(saved.get("generate_denoise_strength", 0.45) or 0.45)) self.prompt_weighting.setChecked(bool(saved.get("generate_prompt_weighting", True))) self._lora_aspect_changed(self.aspect.currentText()) def _lora_aspect_changed(self, aspect: str) -> None: if not self._current_tool() or self._current_tool().id != "lora_generator": return sizes = { "1:1 (Square)": (1024, 1024), "4:3 (Landscape)": (1152, 864), "3:4 (Portrait)": (864, 1152), "3:2 (Landscape)": (1216, 832), "2:3 (Portrait)": (832, 1216), "16:9 (Widescreen)": (1344, 768), "9:16 (Vertical)": (768, 1344), } if aspect in sizes: self.width.setValue(sizes[aspect][0]) self.height.setValue(sizes[aspect][1]) def _browse_reference(self) -> None: path, _ = QFileDialog.getOpenFileName(self, "Choose reference image", self.reference_image.text(), "Images (*.png *.jpg *.jpeg *.webp *.bmp)") if path: self.reference_image.setText(path) def _browse_ddpm_reference(self) -> None: path, _ = QFileDialog.getOpenFileName( self, "Choose DDPM reference image", self.ddpm_reference_image.text(), "Images (*.png *.jpg *.jpeg *.webp *.bmp)", ) if path: self.ddpm_reference_image.setText(path) def _ddpm_custom_size_changed(self, enabled: bool) -> None: self.ddpm_width.setEnabled(enabled) self.ddpm_height.setEnabled(enabled) if self._current_tool() and self._current_tool().id == "ddpm_generator": self.aspect.setEnabled(not enabled) @staticmethod def _model_is_ready(asset: Asset) -> bool: path = Path(asset.path) if asset.trainer == "ddpm": return path.is_dir() and (path / "model_index.json").is_file() if asset.trainer == "flow": return ( path.is_dir() and (path / "flow_model_info.json").is_file() and (path / "unet" / "config.json").is_file() ) if asset.trainer == "lora": return ( path.is_file() and path.suffix.casefold() == ".safetensors" and "_comfy" not in path.stem.casefold() ) or ( path.is_dir() and any( item.is_file() and "_comfy" not in item.stem.casefold() for item in path.glob("*.safetensors") ) ) return path.exists() def _model_changed(self) -> None: tool = self._current_tool() path = self.model.currentData() if not tool: self.model_hint.setText("No registered image generators are available.") elif not path: self.model_hint.setText( f"{tool.name} is ready, but no compatible completed model was found." ) else: self.model_hint.setText(f"Completed pipeline: {path}") self.generate_button.setEnabled(bool(tool and path)) def _apply_preset(self) -> None: name = self.preset.currentText() tool = self._current_tool() if tool and tool.id == "flow_generator": values = { "Balanced": (20, "Heun"), "Fast preview": (8, "Euler"), "Detailed": (30, "Heun"), "Classic": (20, "Euler"), } fallback = (20, "Heun") elif tool and tool.id == "lora_generator": values = { "Balanced": (30, "DPM++ 2M"), "Fast preview": (12, "Euler a"), "Detailed": (50, "DPM++ 2M"), "Classic": (30, "DDIM"), } fallback = (30, "DPM++ 2M") else: values = { "Balanced": (50, "DDIM"), "Fast preview": (20, "DDIM"), "Detailed": (100, "DDIM"), "Classic": (100, "DDPM"), } fallback = (50, "DDIM") steps, sampler = values.get(name, fallback) self.steps.setValue(max(self.steps.minimum(), min(steps, self.steps.maximum()))) if self.sampler.findText(sampler) >= 0: self.sampler.setCurrentText(sampler) def _generate(self) -> None: tool = self._current_tool() model_path = str(self.model.currentData() or "") if tool is None or not model_path: QMessageBox.information( self, "Choose a model", "Choose a registered generator and completed model." ) return raw_seed = self.seed.text().strip() try: seed = int(raw_seed) if raw_seed else 0 except ValueError: QMessageBox.warning(self, "Check the seed", "Seed must be a whole number or left blank for random.") return if seed < 0 or seed > 2_147_483_647: QMessageBox.warning(self, "Check the seed", "Seed must be between 0 and 2,147,483,647.") return extra_arguments = {} if tool.id == "lora_generator": extra_arguments = { "negative_prompt": self.negative_prompt.toPlainText(), "base_model_path": str(self.base_model.currentData() or self.base_model.currentText()), "width": self.width.value(), "height": self.height.value(), "cfg_scale": self.cfg_scale.value(), "lora_strength": self.lora_strength.value(), "reference_image": self.reference_image.text().strip(), "denoise_strength": self.denoise_strength.value(), "prompt_weighting": self.prompt_weighting.isChecked(), "preview_interval": self.preview_interval.value(), } elif tool.id == "ddpm_generator": extra_arguments = { "reference_image": self.ddpm_reference_image.text().strip(), "reference_strength": self.ddpm_reference_strength.value(), "width": self.ddpm_width.value() if self.ddpm_custom_size.isChecked() else 0, "height": self.ddpm_height.value() if self.ddpm_custom_size.isChecked() else 0, "preview_interval": self.preview_interval.value(), } elif tool.id == "flow_generator": extra_arguments = {"preview_interval": self.preview_interval.value()} plan = build_generation_plan( tool, model_name=self.model.currentText(), model_path=model_path, prompt=self.prompt.toPlainText(), image_count=self.images.value(), steps=self.steps.value(), seed=seed, sampler=self.sampler.currentText(), aspect_ratio=self.aspect.currentText(), extra_arguments=extra_arguments, ) job = self.jobs.submit(plan) self._save_generation_settings() self.status.setText( f"Generation job {job.id} queued. You can follow it here or in Jobs / History." ) def _open_generation_cycle(self) -> None: self.assets.discover(self.config) tools = generation_tools(self.registry) trainer_labels = {"ddpm": "DDPM", "flow": "Flow Matching", "lora": "LoRA"} entries = [] for asset in self.assets.assets: if asset.kind != "model" or not self._model_is_ready(asset): continue tool = next((item for item in tools if asset.trainer in item.model_trainers), None) if tool: entries.append({ "name": asset.name, "path": asset.path, "trainer": asset.trainer, "trainer_label": trainer_labels.get(asset.trainer, asset.trainer.title()), "tool": tool, }) if not entries: QMessageBox.information(self, "No completed models", "Complete at least one compatible model before creating a generation cycle.") return dialog = GenerationCycleDialog(entries, self) if dialog.exec() != QDialog.Accepted: return selected = dialog.selected_entries() plans = [] base_seed = dialog.seed.value() for index, entry in enumerate(selected): tool = entry["tool"] options = tool.generation_options samplers = [str(value) for value in options.get("samplers", [])] or ["DDIM"] aspects = [str(value) for value in options.get("aspect_ratios", [])] or ["1:1 (Square)"] extra = {} if tool.id == "lora_generator": saved = self._lora_settings() base_path = str(saved.get("generate_model") or saved.get("last_model") or "") if not base_path: base_asset = next((asset for asset in self.assets.assets if asset.kind == "base_model"), None) base_path = base_asset.path if base_asset else "" if not base_path or not Path(base_path).exists(): QMessageBox.warning( self, "LoRA base model required", f"{entry['name']} needs a valid Stable Diffusion base model before it can join the cycle. " "Choose one in the LoRA generation settings first.", ) return extra = { "base_model_path": base_path, "negative_prompt": str(saved.get("generate_negative_prompt", "")), "width": 1024, "height": 1024, "cfg_scale": float(saved.get("generate_cfg", 7.0) or 7.0), "lora_strength": float(saved.get("generate_lora_strength", 1.0) or 1.0), "denoise_strength": float(saved.get("generate_denoise_strength", 0.45) or 0.45), "prompt_weighting": bool(saved.get("generate_prompt_weighting", True)), } plans.append(build_generation_plan( tool, model_name=entry["name"], model_path=entry["path"], prompt=dialog.prompt.text().strip() or (entry["name"] if tool.id == "lora_generator" else ""), image_count=dialog.image_count.value(), steps=int(options.get("step_default", 30) or 30), seed=base_seed + index * dialog.image_count.value(), sampler=samplers[0], aspect_ratio=aspects[0], extra_arguments=extra, )) plan = combine_generation_plans( plans, display_seconds=dialog.display_seconds.value(), show_labels=dialog.show_labels.isChecked(), loop=dialog.loop.isChecked(), ) job = self.jobs.submit(plan) self._cycle_jobs[job.id] = { "models": [(entry["name"], entry["trainer_label"]) for entry in selected], "display_seconds": dialog.display_seconds.value(), "show_labels": dialog.show_labels.isChecked(), "loop": dialog.loop.isChecked(), "fullscreen": dialog.fullscreen.isChecked(), } self.status.setText(f"Generation cycle {job.id} queued with {len(selected)} models.") def _play_finished_cycle(self, job: Job, settings: dict) -> None: records = [] history_root = self.root / "data" / "generations" if history_root.is_dir(): for metadata in history_root.rglob(f"*{job.id}*.json"): record = GenerationRecord.from_metadata(metadata) if record: records.append(record) model_order = {name: index for index, (name, _trainer) in enumerate(settings["models"])} records.sort(key=lambda record: model_order.get(record.model_name, len(model_order))) trainer_by_model = dict(settings["models"]) slides = [ (path, f"{record.model_name} — {trainer_by_model.get(record.model_name, record.provider_name)}") for record in records for path in record.images ] if not slides: self.status.setText(f"Generation cycle {job.id} finished, but no cycle images were found.") return player = CycleSlideshowDialog( slides, display_seconds=settings["display_seconds"], show_labels=settings["show_labels"], loop=settings["loop"], parent=self, ) if settings.get("fullscreen"): player.showFullScreen() player.exec() def _save_generation_settings(self) -> None: self.config.update({ "generation_settings": { "provider": self.provider.currentData(), "model_path": self.model.currentData(), "preset": self.preset.currentText(), "images": self.images.value(), "steps": self.steps.value(), "preview_interval": self.preview_interval.value(), "sampler": self.sampler.currentText(), "aspect_ratio": self.aspect.currentText(), "seed": self.seed.text(), "prompt": self.prompt.toPlainText(), "base_model_path": self.base_model.currentData() or self.base_model.currentText(), "negative_prompt": self.negative_prompt.toPlainText(), "cfg_scale": self.cfg_scale.value(), "lora_strength": self.lora_strength.value(), "denoise_strength": self.denoise_strength.value(), "prompt_weighting": self.prompt_weighting.isChecked(), "ddpm_reference_image": self.ddpm_reference_image.text(), "ddpm_reference_strength": self.ddpm_reference_strength.value(), "ddpm_custom_size": self.ddpm_custom_size.isChecked(), "ddpm_width": self.ddpm_width.value(), "ddpm_height": self.ddpm_height.value(), } }) def _connect_generation_setting_persistence(self) -> None: self.provider.currentIndexChanged.connect(self._save_generation_settings) self.model.currentIndexChanged.connect(self._save_generation_settings) self.preset.currentTextChanged.connect(self._save_generation_settings) self.images.valueChanged.connect(self._save_generation_settings) self.steps.valueChanged.connect(self._save_generation_settings) self.preview_interval.valueChanged.connect(self._save_generation_settings) self.sampler.currentTextChanged.connect(self._save_generation_settings) self.aspect.currentTextChanged.connect(self._save_generation_settings) self.seed.editingFinished.connect(self._save_generation_settings) self.prompt.textChanged.connect(self._save_generation_settings) self.base_model.currentIndexChanged.connect(self._save_generation_settings) self.base_model.editTextChanged.connect(self._save_generation_settings) self.negative_prompt.textChanged.connect(self._save_generation_settings) self.cfg_scale.valueChanged.connect(self._save_generation_settings) self.lora_strength.valueChanged.connect(self._save_generation_settings) self.denoise_strength.valueChanged.connect(self._save_generation_settings) self.prompt_weighting.toggled.connect(self._save_generation_settings) self.ddpm_reference_image.editingFinished.connect(self._save_generation_settings) self.ddpm_reference_strength.valueChanged.connect(self._save_generation_settings) self.ddpm_custom_size.toggled.connect(self._save_generation_settings) self.ddpm_width.valueChanged.connect(self._save_generation_settings) self.ddpm_height.valueChanged.connect(self._save_generation_settings) def _restore_generation_settings(self) -> None: saved = self.config.get("generation_settings", {}) if not isinstance(saved, dict): return provider_index = self.provider.findData(saved.get("provider")) if provider_index >= 0: self.provider.setCurrentIndex(provider_index) model_index = self.model.findData(saved.get("model_path")) if model_index >= 0: self.model.setCurrentIndex(model_index) for combo, key in ((self.preset, "preset"), (self.sampler, "sampler"), (self.aspect, "aspect_ratio")): value = str(saved.get(key, "")) if combo.findText(value) >= 0: combo.setCurrentText(value) for spin, key in ((self.images, "images"), (self.steps, "steps"), (self.preview_interval, "preview_interval"), (self.ddpm_reference_strength, "ddpm_reference_strength"), (self.ddpm_width, "ddpm_width"), (self.ddpm_height, "ddpm_height")): try: spin.setValue(max(spin.minimum(), min(int(saved.get(key, spin.value())), spin.maximum()))) except (TypeError, ValueError): pass self.seed.setText(str(saved.get("seed", ""))) self.prompt.setPlainText(str(saved.get("prompt", ""))) base_model_path = str(saved.get("base_model_path", "")) base_index = self.base_model.findData(base_model_path) if base_index >= 0: self.base_model.setCurrentIndex(base_index) elif base_model_path: self.base_model.setCurrentText(base_model_path) self.negative_prompt.setPlainText(str(saved.get("negative_prompt", self.negative_prompt.toPlainText()))) try: self.cfg_scale.setValue(float(saved.get("cfg_scale", self.cfg_scale.value()))) self.lora_strength.setValue(float(saved.get("lora_strength", self.lora_strength.value()))) self.denoise_strength.setValue(float(saved.get("denoise_strength", self.denoise_strength.value()))) except (TypeError, ValueError): pass self.prompt_weighting.setChecked(bool(saved.get("prompt_weighting", self.prompt_weighting.isChecked()))) self.ddpm_reference_image.setText(str(saved.get("ddpm_reference_image", ""))) self.ddpm_custom_size.setChecked(bool(saved.get("ddpm_custom_size", False))) def _load_history(self) -> None: selected_image = None current = self.gallery.currentItem() if current: selected_image = current.data(Qt.UserRole) self.records = load_generation_history(self.root) self.gallery.clear() image_total = 0 selected_row = -1 for record_index, record in enumerate(self.records): for image_index, path in enumerate(record.images): if str(path) in self.hidden_history_images: continue image_total += 1 seed = record.seed + image_index item = QListWidgetItem( QIcon(_thumbnail(path, 170, 128)), f"{record.model_name}\nSeed {seed}", ) payload = { "record": record_index, "image": str(path), "seed": seed, } item.setData(Qt.UserRole, payload) self.gallery.addItem(item) if selected_image and selected_image.get("image") == str(path): selected_row = self.gallery.count() - 1 self.history_summary.setText( f"{len(self.records)} batches · {image_total} images" if image_total else "No generations yet" ) self.clear_history_button.setEnabled(image_total > 0) if selected_row >= 0: self.gallery.setCurrentRow(selected_row) elif self.gallery.count(): self.gallery.setCurrentRow(0) else: self._selection_changed(None, None) def _clear_displayed_history(self) -> None: for index in range(self.gallery.count()): payload = self.gallery.item(index).data(Qt.UserRole) if isinstance(payload, dict) and payload.get("image"): self.hidden_history_images.add(str(payload["image"])) self.gallery.clear() self.history_summary.setText("No images displayed") self.clear_history_button.setEnabled(False) self._selection_changed(None, None) self.status.setText("Displayed generation history cleared. Image files were not deleted.") def _selection(self) -> tuple[GenerationRecord, Path, int] | None: item = self.gallery.currentItem() if not item: return None payload = item.data(Qt.UserRole) try: return ( self.records[int(payload["record"])], Path(str(payload["image"])), int(payload["seed"]), ) except (KeyError, IndexError, TypeError, ValueError): return None def _selection_changed(self, _current=None, _previous=None) -> None: selection = self._selection() enabled = selection is not None self.open_image_button.setEnabled(enabled) self.open_folder_button.setEnabled(enabled) self.reuse_button.setEnabled(enabled) if not selection: self.detail.setText("Select an image to see its reproducibility settings.") return record, _path, seed = selection note = f" · Note: {record.prompt}" if record.prompt else "" self.detail.setText( f"{record.provider_name} · {record.model_name} · Seed {seed} · " f"{record.steps} steps · {record.sampler} · {record.aspect_ratio}{note}" ) def _open_image(self) -> None: selection = self._selection() if selection: QDesktopServices.openUrl(QUrl.fromLocalFile(str(selection[1]))) def _open_folder(self) -> None: selection = self._selection() if selection: QDesktopServices.openUrl(QUrl.fromLocalFile(str(selection[0].folder))) def _reuse_settings(self) -> None: selection = self._selection() if not selection: return record, _path, seed = selection provider_index = self.provider.findData(record.provider_id) if provider_index >= 0: self.provider.setCurrentIndex(provider_index) model_index = self.model.findData(record.model_path) if model_index >= 0: self.model.setCurrentIndex(model_index) self.prompt.setPlainText(record.prompt) self.seed.setText(str(seed)) self.steps.setValue(max(self.steps.minimum(), min(record.steps, self.steps.maximum()))) self.sampler.setCurrentText(record.sampler) self.aspect.setCurrentText(record.aspect_ratio) self.status.setText("Settings restored from generation history.") def _job_updated(self, job: Job) -> None: generation_ids = {tool.id for tool in generation_tools(self.registry)} if not any(step.tool_id in generation_ids for step in job.plan.steps): return if job.status == JobStatus.FINISHED: self.status.setText(f"Generation job {job.id} finished successfully.") self._load_history() cycle = self._cycle_jobs.pop(job.id, None) if cycle: QTimer.singleShot(0, lambda current=job, settings=cycle: self._play_finished_cycle(current, settings)) elif job.status == JobStatus.FAILED: self.status.setText(f"Generation job {job.id} failed: {job.error or 'Unknown error'}") elif job.status == JobStatus.CANCELLED: self.status.setText(f"Generation job {job.id} was cancelled.") elif job.status in {JobStatus.RUNNING, JobStatus.PAUSED}: self.status.setText(f"Generation job {job.id}: {job.status.value} · {job.progress}%")