Add support for Krita Style & Prompt node (#2276)

Features:
- Prompt and Style widgets sync across Generate/Live/Animation and now Graph workspaces
- Full prompt evaluation: wildcards, comments, layer extraction,
  style merging, LoRA extraction
- LoRAs applied to model output
- Error shown if workflow contains multiple KritaPromptStyle nodes
  (only one allowed since prompts sync to shared fields)

Author: leha-games

ai_diffusion/api.py

@@ -169,6 +169,10 @@ class UpscaleInput:
 class CustomWorkflowInput:
     workflow: dict
     params: dict[str, Any]
+    positive_evaluated: str = ""
+    negative_evaluated: str = ""
+    models: CheckpointInput | None = None
+    sampling: SamplingInput | None = None
 
 
 @dataclass

ai_diffusion/custom_workflow.py

@@ -13,6 +13,7 @@
 from .api import WorkflowInput, InpaintContext
 from .client import OutputBatchMode, TextOutput, ClientOutput, JobInfoOutput
 from .comfy_workflow import ComfyWorkflow, ComfyNode
+from .localization import translate as _
 from .connection import Connection, ConnectionState
 from .image import Bounds, Image, Mask
 from .jobs import Job, JobParams, JobQueue, JobKind
@@ -373,6 +374,7 @@ class CustomWorkspace(QObject, ObservableProperties):
     has_result = Property(False)
     outputs = Property({})
     params_ui_height = Property(100, persist=True)
+    validation_error = Property("")
 
     workflow_id_changed = pyqtSignal(str)
     graph_changed = pyqtSignal()
@@ -383,6 +385,7 @@ class CustomWorkspace(QObject, ObservableProperties):
     has_result_changed = pyqtSignal(bool)
     outputs_changed = pyqtSignal(dict)
     params_ui_height_changed = pyqtSignal(int)
+    validation_error_changed = pyqtSignal(str)
     modified = pyqtSignal(QObject, str)
 
     _live_poll_rate = 0.1
@@ -419,10 +422,20 @@ def _update_workflow(self, idx: QModelIndex, _: QModelIndex):
         if wf.id == self._workflow_id:
             self._workflow = wf
             self._graph = self._workflow.workflow
+            self._validate_workflow(self._graph)
             self._metadata = list(workflow_parameters(self._graph))
             self.params = _coerce(self.params, self._metadata)
             self.graph_changed.emit()
 
+    def _validate_workflow(self, wf: ComfyWorkflow):
+        style_and_prompt_node_count = sum(1 for _ in wf.find(type="ETN_KritaStyleAndPrompt"))
+        if style_and_prompt_node_count > 1:
+            self.validation_error = _(
+                "Workflow contains multiple `Krita Style & Prompt` nodes. Only one is allowed since prompts sync across workspaces."
+            )
+        else:
+            self.validation_error = ""
+
     def _set_workflow_id(self, id: str):
         if self._workflow_id == id:
             return

ai_diffusion/model.py

@@ -19,14 +19,19 @@
 from .api import FillMode, ImageInput, CustomWorkflowInput, UpscaleInput
 from .api import InpaintMode, InpaintContext, InpaintParams
 from .localization import translate as _
-from .util import clamp, ensure, trim_text, client_logger as log
+from .util import clamp, ensure, unique, trim_text, client_logger as log
 from .settings import ApplyBehavior, ApplyRegionBehavior, GenerationFinishedAction, ImageFileFormat
 from .settings import settings
 from .network import NetworkError
 from .image import Extent, Image, Mask, Bounds, DummyImage
 from .client import Client, ClientMessage, ClientEvent, ClientOutput
 from .client import is_style_supported, filter_supported_styles, resolve_arch
-from .custom_workflow import CustomWorkspace, WorkflowCollection, CustomGenerationMode
+from .custom_workflow import (
+    CustomWorkspace,
+    WorkflowCollection,
+    CustomGenerationMode,
+    ComfyWorkflow,
+)
 from .document import Document, KritaDocument, SelectionModifiers
 from .layer import Layer, LayerType, RestoreActiveLayer
 from .pose import Pose
@@ -149,6 +154,7 @@ def __init__(self, document: Document, connection: Connection, workflows: Workfl
         self.jobs.selection_changed.connect(self.update_preview)
         connection.state_changed.connect(self._init_on_connect)
         connection.error_changed.connect(self._forward_error)
+        self.custom.validation_error_changed.connect(self._forward_validation_error)
         Styles.list().changed.connect(self._init_on_connect)
         self._init_on_connect()
 
@@ -165,6 +171,12 @@ def _init_on_connect(self):
     def _forward_error(self, error: str):
         self.report_error(error if error else no_error)
 
+    def _forward_validation_error(self, error: str):
+        if error:
+            self.report_error(Error(ErrorKind.warning, error))
+        else:
+            self.clear_error()
+
     def generate(self):
         """Enqueue image generation for the current setup."""
         self._generate(self.queue_mode)
@@ -488,14 +500,29 @@ async def _generate_custom(self, previous_input: WorkflowInput | None):
             img_input.hires_mask = mask.to_image(bounds.extent) if mask else None
 
             params = self.custom.collect_parameters(self.layers, canvas_bounds, is_anim)
+
+            has_synced_style_and_prompt = (
+                next(wf.find(type="ETN_KritaStyleAndPrompt"), None) is not None
+            )
+            custom_input = CustomWorkflowInput(wf.root, params)
+            prompt_meta = {}
+            if has_synced_style_and_prompt:
+                custom_input, prompt_meta = self._prepare_synced_style_and_prompt(
+                    params, seed, custom_input, wf
+                )
+
             input = WorkflowInput(
                 WorkflowKind.custom,
                 img_input,
                 sampling=SamplingInput("custom", "custom", 1, 1000, seed=seed),
                 inpaint=InpaintParams(InpaintMode.fill, bounds),
-                custom_workflow=CustomWorkflowInput(wf.root, params),
+                custom_workflow=custom_input,
             )
-            job_params = JobParams(bounds, self.custom.job_name, metadata=self.custom.params)
+
+            metadata: dict[str, Any] = dict(self.custom.params)
+            metadata.update(prompt_meta)
+
+            job_params = JobParams(bounds, self.custom.job_name, metadata=metadata)
             job_kind = {
                 CustomGenerationMode.regular: JobKind.diffusion,
                 CustomGenerationMode.live: JobKind.live_preview,
@@ -513,6 +540,46 @@ async def _generate_custom(self, previous_input: WorkflowInput | None):
             self.report_error(util.log_error(e))
             return False
 
+    def _prepare_synced_style_and_prompt(
+        self,
+        params: dict[str, Any],
+        seed: int,
+        custom_input: CustomWorkflowInput,
+        wf: ComfyWorkflow,
+    ) -> tuple[CustomWorkflowInput, dict[str, Any]]:
+        """Prepare prompts and models for ETN_KritaStyleAndPrompt node.
+        Returns updated CustomWorkflowInput with evaluated prompts, models, sampling, and metadata for job history.
+        """
+        style = self.style
+
+        style_node = next(wf.find(type="ETN_KritaStyleAndPrompt"), None)
+        is_live = style_node.input("sampler_preset", "auto") == "live" if style_node else False
+
+        checkpoint_input = style.get_models(self._connection.client.models.checkpoints)
+        sampling = workflow._sampling_from_style(style, 1.0, is_live)
+
+        positive = self.regions.positive
+        negative = self.regions.negative
+
+        cond = ConditioningInput(positive, negative)
+        arch = resolve_arch(style, self._connection.client_if_connected)
+        prepared = workflow.prepare_prompts(cond, style, seed, arch, FileLibrary.instance())
+
+        merged_loras = unique(checkpoint_input.loras + prepared.loras, key=lambda l: l.name)
+        checkpoint_input.loras = merged_loras
+
+        custom_input = replace(
+            custom_input,
+            positive_evaluated=prepared.metadata["prompt_final"],
+            negative_evaluated=prepared.metadata["negative_prompt_final"],
+            models=checkpoint_input,
+            sampling=sampling,
+        )
+
+        meta = dict(prepared.metadata)
+        meta["style"] = style.filename
+        return custom_input, meta
+
     def _get_current_image(self, bounds: Bounds):
         exclude = []
         if self.workspace is not Workspace.live:

ai_diffusion/ui/custom_workflow.py

@@ -25,6 +25,7 @@
 from .live import LivePreviewArea
 from .switch import SwitchWidget
 from .widget import TextPromptWidget, WorkspaceSelectWidget, StyleSelectWidget, ErrorBox
+from .region import ActiveRegionWidget, PromptHeader
 from .settings_widgets import ExpanderButton
 from . import theme
 from .theme import SignalBlocker
@@ -700,6 +701,16 @@ def __init__(self):
         self._params_scroll.setFrameShape(QFrame.Shape.NoFrame)
         self._params_scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
 
+        self._style_widget = StyleSelectWidget(self)
+        self._style_widget.setVisible(False)  # Hidden until workflow has synced style
+
+        self._prompt_widget = ActiveRegionWidget(
+            self._model.regions, self, header=PromptHeader.none
+        )
+        self._prompt_widget.setVisible(False)  # Hidden until workflow has synced prompts
+        self._prompt_widget.positive.activated.connect(self._generate)
+        self._prompt_widget.negative.activated.connect(self._generate)
+
         self._bottom = QWidget(self)
 
         self._generate_button = GenerateButton(JobKind.diffusion, self._bottom)
@@ -771,6 +782,8 @@ def __init__(self):
         header_layout.addWidget(self._workflow_select_widgets)
         header_layout.addWidget(self._workflow_edit_widgets)
         layout.addLayout(header_layout)
+        layout.addWidget(self._style_widget)
+        layout.addWidget(self._prompt_widget)
         layout.addWidget(self._splitter)
         actions_layout = QHBoxLayout()
         actions_layout.setSpacing(0)
@@ -808,6 +821,7 @@ def model(self, model: Model):
             self._model = model
             self._model_bindings = [
                 bind(model, "workspace", self._workspace_select, "value", Bind.one_way),
+                bind(model, "style", self._style_widget, "value"),
                 bind(model, "error", self._error_box, "error", Bind.one_way),
                 bind_combo(model.custom, "workflow_id", self._workflow_select, Bind.one_way),
                 bind(model.custom, "outputs", self._outputs, "value", Bind.one_way),
@@ -822,6 +836,7 @@ def model(self, model: Model):
             self._queue_button.model = model
             self._progress_bar.model = model
             self._history.model_ = model
+            self._prompt_widget.region = model.regions
             self._update_current_workflow()
             self._update_ui()
             self._set_params_height(model.custom.params_ui_height)
@@ -874,12 +889,21 @@ def _update_current_workflow(self):
         if not self.model.custom.workflow:
             self._save_workflow_button.setEnabled(False)
             self._delete_workflow_button.setEnabled(False)
+            self._style_widget.setVisible(False)
+            self._prompt_widget.setVisible(False)
             return
         self._save_workflow_button.setEnabled(True)
         self._delete_workflow_button.setEnabled(
             self.model.custom.workflow.source is WorkflowSource.local
         )
 
+        graph = self.model.custom.graph
+        has_synced_style_and_prompt = (
+            graph is not None and next(graph.find(type="ETN_KritaStyleAndPrompt"), None) is not None
+        )
+        self._style_widget.setVisible(has_synced_style_and_prompt)
+        self._prompt_widget.setVisible(has_synced_style_and_prompt)
+
         if self._params_widget:
             self._params_scroll.setWidget(None)
             self._params_widget.deleteLater()

ai_diffusion/workflow.py

@@ -1420,6 +1420,20 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None =
                 outputs[node.output(6)] = sampling.scheduler
                 outputs[node.output(7)] = sampling.total_steps
                 outputs[node.output(8)] = sampling.cfg_scale
+
+            case "ETN_KritaStyleAndPrompt":
+                checkpoint_input = ensure(input.models)
+                sampling = ensure(input.sampling)
+                model, clip, vae = load_checkpoint_with_lora(w, checkpoint_input, models)
+                outputs[node.output(0)] = model
+                outputs[node.output(1)] = clip.model
+                outputs[node.output(2)] = vae
+                outputs[node.output(3)] = input.positive_evaluated
+                outputs[node.output(4)] = input.negative_evaluated
+                outputs[node.output(5)] = sampling.sampler
+                outputs[node.output(6)] = sampling.scheduler
+                outputs[node.output(7)] = sampling.total_steps
+                outputs[node.output(8)] = sampling.cfg_scale
             case _:
                 mapped_inputs = {k: map_input(v) for k, v in node.inputs.items()}
                 mapped = ComfyNode(node.id, node.type, mapped_inputs)