griptape-nodes 0.38.1__tar.gz → 0.40.0__tar.gz

{griptape_nodes-0.38.1 → griptape_nodes-0.40.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: griptape-nodes
-Version: 0.38.1
+Version: 0.40.0
 Summary: Add your description here
 License-File: LICENSE
 Requires-Python: ~=3.12

griptape_nodes-0.40.0/libraries/griptape_nodes_advanced_media_library/diffusers_nodes_library/common/utils/huggingface_utils.py

@@ -0,0 +1,126 @@
+import logging
+from functools import cache
+from pathlib import Path
+from typing import Any
+
+from huggingface_hub import scan_cache_dir
+from huggingface_hub.constants import HF_HUB_CACHE
+
+logger = logging.getLogger("griptape_nodes")
+
+
+def list_repo_revisions_in_cache(repo_id: str) -> list[tuple[str, str]]:
+    """Returns a list of (repo_id, revision) tuples matching repo_id in the huggingface cache."""
+    # Use quick scan for diffuser repos, fallback to scan_cache_dir only on errors
+    try:
+        repos = quick_scan_diffuser_repos(HF_HUB_CACHE)
+        results = [(repo["name"], repo["hash"]) for repo in repos if repo["name"] == repo_id]
+    except Exception:
+        logger.exception("Failed to quick scan diffuser repos, falling back to scan_cache_dir.")
+    else:
+        return results
+
+    # Fallback to original implementation
+    cache_info = scan_cache_dir()
+    results = []
+    for repo in cache_info.repos:
+        if repo.repo_id == repo_id:
+            for revision in repo.revisions:
+                results.append((repo.repo_id, revision.commit_hash))
+    return results
+
+
+def list_repo_revisions_with_file_in_cache(repo_id: str, file: str) -> list[tuple[str, str]]:
+    """Returns a list of (repo_id, revision) tuples matching repo_id in the huggingface cache if it contains file."""
+    # Use quick scan for diffuser repos, check if file exists
+    try:
+        repos = quick_scan_diffuser_repos(HF_HUB_CACHE)
+        results = [
+            (repo["name"], repo["hash"])
+            for repo in repos
+            if repo["name"] == repo_id and (Path(repo["path"]) / file).exists()
+        ]
+    except Exception:
+        logger.exception("Failed to quick scan diffuser repos, falling back to scan_cache_dir.")
+    else:
+        return results
+
+    # Fallback to original implementation
+    cache_info = scan_cache_dir()
+    results = []
+    for repo in cache_info.repos:
+        if repo.repo_id == repo_id:
+            for revision in repo.revisions:
+                if any(f.file_name == file for f in revision.files):
+                    results.append((repo.repo_id, revision.commit_hash))
+    return results
+
+
+class ModelCache:
+    @cache  # noqa: B019
+    def from_pretrained(self, cls: Any, *args, **kwargs) -> Any:
+        return cls.from_pretrained(*args, **kwargs)
+
+
+model_cache = ModelCache()
+
+
+def quick_scan_diffuser_repos(cache_dir: str) -> list[dict[str, Any]]:
+    """Quick scan of diffuser repositories in huggingface cache directory.
+
+    Adapted from https://github.com/huggingface/huggingface_hub/issues/1564#issuecomment-1710764361
+
+    Args:
+        cache_dir: Path to huggingface cache directory
+
+    Returns:
+        List of dictionaries containing repo info with keys:
+        - name: Repository name
+        - filename: Repository name (alias for name)
+        - path: Path to the snapshot folder
+        - hash: Commit hash
+        - mtime: Modification time
+        - model_info: Path to model_info.json file
+    """
+    diffuser_repos = []
+    cache_path = Path(cache_dir)
+
+    if not cache_path.exists():
+        return diffuser_repos
+
+    for folder_path in cache_path.iterdir():
+        folder = folder_path.name
+        if not folder_path.is_dir() or "--" not in folder:
+            continue
+        _, name = folder.split("--", maxsplit=1)
+        name = name.replace("--", "/")
+
+        snapshots_dir = folder_path / "snapshots"
+        if not snapshots_dir.exists():
+            continue
+
+        snapshots = [p.name for p in snapshots_dir.iterdir() if p.is_dir()]
+        if len(snapshots) == 0:
+            continue
+
+        commit = snapshots[-1]
+        snapshot_path = snapshots_dir / commit
+
+        if (snapshot_path / "hidden").exists():
+            continue
+
+        mtime = snapshot_path.stat().st_mtime
+        info = snapshot_path / "model_info.json"
+
+        diffuser_repos.append(
+            {
+                "name": name,
+                "filename": name,
+                "path": str(snapshot_path),
+                "hash": commit,
+                "mtime": mtime,
+                "model_info": str(info),
+            }
+        )
+
+    return diffuser_repos

griptape_nodes-0.40.0/libraries/griptape_nodes_advanced_media_library/diffusers_nodes_library/pipelines/flux/flux_kontext_pipeline.py

@@ -0,0 +1,95 @@
+import logging
+from typing import Any
+
+import diffusers  # type: ignore[reportMissingImports]
+import torch  # type: ignore[reportMissingImports]
+
+from diffusers_nodes_library.common.parameters.log_parameter import (  # type: ignore[reportMissingImports]
+    LogParameter,  # type: ignore[reportMissingImports]
+)
+from diffusers_nodes_library.common.utils.huggingface_utils import model_cache  # type: ignore[reportMissingImports]
+from diffusers_nodes_library.pipelines.flux.flux_kontext_pipeline_memory_footprint import (
+    optimize_flux_kontext_pipeline_memory_footprint,
+)  # type: ignore[reportMissingImports]
+from diffusers_nodes_library.pipelines.flux.flux_kontext_pipeline_parameters import (
+    FluxKontextPipelineParameters,  # type: ignore[reportMissingImports]
+)
+from diffusers_nodes_library.pipelines.flux.flux_loras_parameter import FluxLorasParameter
+from griptape_nodes.exe_types.core_types import Parameter
+from griptape_nodes.exe_types.node_types import AsyncResult, ControlNode
+
+logger = logging.getLogger("diffusers_nodes_library")
+
+
+class FluxKontextPipeline(ControlNode):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.pipe_params = FluxKontextPipelineParameters(self)
+        self.loras_params = FluxLorasParameter(self)
+        self.log_params = LogParameter(self)
+        self.pipe_params.add_input_parameters()
+        self.loras_params.add_input_parameters()
+        self.pipe_params.add_output_parameters()
+        self.log_params.add_output_parameters()
+
+    def after_value_set(self, parameter: Parameter, value: Any) -> None:
+        self.pipe_params.after_value_set(parameter, value)
+
+    def validate_before_node_run(self) -> list[Exception] | None:
+        errors = self.pipe_params.validate_before_node_run()
+        return errors or None
+
+    def validate_before_workflow_run(self) -> list[Exception] | None:
+        return self.pipe_params.validate_before_node_run()
+
+    def preprocess(self) -> None:
+        self.pipe_params.preprocess()
+
+    def process(self) -> AsyncResult | None:
+        yield lambda: self._process()
+
+    def _process(self) -> AsyncResult | None:
+        self.preprocess()
+        self.pipe_params.publish_output_image_preview_placeholder()
+        self.log_params.append_to_logs("Preparing models...\n")
+
+        with self.log_params.append_profile_to_logs("Loading model metadata"):
+            base_repo_id, base_revision = self.pipe_params.get_repo_revision()
+            pipe = model_cache.from_pretrained(
+                diffusers.FluxKontextPipeline,
+                pretrained_model_name_or_path=base_repo_id,
+                revision=base_revision,
+                torch_dtype=torch.bfloat16,
+                local_files_only=True,
+            )
+
+        with self.log_params.append_profile_to_logs("Loading model"), self.log_params.append_logs_to_logs(logger):
+            optimize_flux_kontext_pipeline_memory_footprint(pipe)
+
+        with (
+            self.log_params.append_profile_to_logs("Configuring flux loras"),
+            self.log_params.append_logs_to_logs(logger),
+        ):
+            self.loras_params.configure_loras(pipe)
+
+        num_inference_steps = self.pipe_params.get_num_inference_steps()
+
+        def callback_on_step_end(
+            pipe: diffusers.FluxKontextPipeline,
+            i: int,
+            _t: Any,
+            callback_kwargs: dict,
+        ) -> dict:
+            if i < num_inference_steps - 1:
+                self.pipe_params.publish_output_image_preview_latents(pipe, callback_kwargs["latents"])
+                self.log_params.append_to_logs(f"Starting inference step {i + 2} of {num_inference_steps}...\n")
+            return {}
+
+        self.log_params.append_to_logs(f"Starting inference step 1 of {num_inference_steps}...\n")
+        output_image_pil = pipe(
+            **self.pipe_params.get_pipe_kwargs(),
+            output_type="pil",
+            callback_on_step_end=callback_on_step_end,
+        ).images[0]
+        self.pipe_params.publish_output_image(output_image_pil)
+        self.log_params.append_to_logs("Done.\n")

griptape_nodes-0.40.0/libraries/griptape_nodes_advanced_media_library/diffusers_nodes_library/pipelines/flux/flux_kontext_pipeline_memory_footprint.py

@@ -0,0 +1,129 @@
+import logging
+from functools import cache
+
+import diffusers  # type: ignore[reportMissingImports]
+import torch  # type: ignore[reportMissingImports]
+
+from diffusers_nodes_library.common.utils.torch_utils import (
+    get_best_device,
+    get_total_memory_footprint,  # type: ignore[reportMissingImports]
+    print_pipeline_memory_footprint,
+    to_human_readable_size,  # type: ignore[reportMissingImports]
+)
+
+logger = logging.getLogger("diffusers_nodes_library")
+
+
+FLUX_KONTEXT_PIPELINE_COMPONENT_NAMES = [
+    "vae",
+    "text_encoder",
+    "text_encoder_2",
+    "transformer",
+]
+
+
+def print_flux_kontext_pipeline_memory_footprint(pipe: diffusers.FluxKontextPipeline) -> None:
+    """Print memory footprint for the main sub-modules of Flux Kontext pipelines."""
+    print_pipeline_memory_footprint(pipe, FLUX_KONTEXT_PIPELINE_COMPONENT_NAMES)
+
+
+def _check_cuda_memory_sufficient(pipe: diffusers.FluxKontextPipeline, device: torch.device) -> bool:
+    """Check if CUDA device has sufficient memory for the pipeline."""
+    model_memory = get_total_memory_footprint(pipe, FLUX_KONTEXT_PIPELINE_COMPONENT_NAMES)
+    total_memory = torch.cuda.get_device_properties(device).total_memory
+    free_memory = total_memory - torch.cuda.memory_allocated(device)
+    return model_memory <= free_memory
+
+
+def _check_mps_memory_sufficient(pipe: diffusers.FluxKontextPipeline) -> bool:
+    """Check if MPS device has sufficient memory for the pipeline."""
+    model_memory = get_total_memory_footprint(pipe, FLUX_KONTEXT_PIPELINE_COMPONENT_NAMES)
+    recommended_max_memory = torch.mps.recommended_max_memory()
+    free_memory = recommended_max_memory - torch.mps.current_allocated_memory()
+    return model_memory <= free_memory
+
+
+def _log_memory_info(pipe: diffusers.FluxKontextPipeline, device: torch.device) -> None:
+    """Log memory information for the device."""
+    model_memory = get_total_memory_footprint(pipe, FLUX_KONTEXT_PIPELINE_COMPONENT_NAMES)
+
+    if device.type == "cuda":
+        total_memory = torch.cuda.get_device_properties(device).total_memory
+        free_memory = total_memory - torch.cuda.memory_allocated(device)
+        logger.info("Total memory on %s: %s", device, to_human_readable_size(total_memory))
+        logger.info("Free memory on %s: %s", device, to_human_readable_size(free_memory))
+    elif device.type == "mps":
+        recommended_max_memory = torch.mps.recommended_max_memory()
+        free_memory = recommended_max_memory - torch.mps.current_allocated_memory()
+        logger.info("Recommended max memory on %s: %s", device, to_human_readable_size(recommended_max_memory))
+        logger.info("Free memory on %s: %s", device, to_human_readable_size(free_memory))
+
+    logger.info("Require memory for Flux Kontext Pipeline: %s", to_human_readable_size(model_memory))
+
+
+@cache
+def optimize_flux_kontext_pipeline_memory_footprint(pipe: diffusers.FluxKontextPipeline) -> None:
+    """Optimize pipeline memory footprint with incremental VRAM checking."""
+    device = get_best_device()
+
+    if device.type == "cuda":
+        _log_memory_info(pipe, device)
+
+        if _check_cuda_memory_sufficient(pipe, device):
+            logger.info("Sufficient memory on %s for Pipeline.", device)
+            logger.info("Moving pipeline to %s", device)
+            pipe.to(device)
+            return
+
+        logger.warning("Insufficient memory on %s for Pipeline. Applying VRAM optimizations.", device)
+
+        # Apply fp8 layerwise caching
+        logger.info("Enabling fp8 layerwise caching for transformer")
+        pipe.transformer.enable_layerwise_casting(
+            storage_dtype=torch.float8_e4m3fn,
+            compute_dtype=torch.bfloat16,
+        )
+
+        if _check_cuda_memory_sufficient(pipe, device):
+            logger.info("Sufficient memory after fp8 optimization. Moving pipeline to %s", device)
+            pipe.to(device)
+            return
+
+        # Apply sequential CPU offload
+        logger.info("Still insufficient memory. Enabling sequential cpu offload")
+        pipe.enable_sequential_cpu_offload()
+
+        if _check_cuda_memory_sufficient(pipe, device):
+            logger.info("Sufficient memory after sequential cpu offload")
+            return
+
+        # Apply VAE slicing as final optimization
+        logger.info("Enabling vae slicing")
+        pipe.enable_vae_slicing()
+
+        # Final check after all optimizations
+        if not _check_cuda_memory_sufficient(pipe, device):
+            logger.warning("Memory may still be insufficient after all optimizations, but will try anyway")
+
+        # Intentionally not calling pipe.to(device) here because sequential_cpu_offload
+        # manages device placement automatically
+
+    elif device.type == "mps":
+        _log_memory_info(pipe, device)
+
+        if _check_mps_memory_sufficient(pipe):
+            logger.info("Sufficient memory on %s for Pipeline.", device)
+            logger.info("Moving pipeline to %s", device)
+            pipe.to(device)
+            return
+
+        logger.warning("Insufficient memory on %s for Pipeline.", device)
+        logger.info("Enabling vae slicing")
+        pipe.enable_vae_slicing()
+
+        # Final check after VAE slicing
+        if not _check_mps_memory_sufficient(pipe):
+            logger.warning("Memory may still be insufficient after optimizations, but will try anyway")
+
+        # Intentionally not calling pipe.to(device) here when memory is insufficient
+        # to avoid potential OOM errors

griptape_nodes-0.40.0/libraries/griptape_nodes_advanced_media_library/diffusers_nodes_library/pipelines/flux/flux_kontext_pipeline_parameters.py

@@ -0,0 +1,303 @@
+import logging
+from typing import Any
+
+import diffusers  # type: ignore[reportMissingImports]
+import PIL.Image
+from griptape.artifacts import ImageUrlArtifact
+from PIL.Image import Image
+from pillow_nodes_library.utils import (  # type: ignore[reportMissingImports]
+    image_artifact_to_pil,
+    pil_to_image_artifact,
+)
+from utils.image_utils import load_image_from_url_artifact
+
+from diffusers_nodes_library.common.parameters.huggingface_repo_parameter import HuggingFaceRepoParameter
+from diffusers_nodes_library.common.parameters.seed_parameter import SeedParameter
+from griptape_nodes.exe_types.core_types import Parameter, ParameterMode
+from griptape_nodes.exe_types.node_types import BaseNode
+from griptape_nodes.traits.options import Options
+
+logger = logging.getLogger("diffusers_nodes_library")
+
+
+class FluxKontextPipelineParameters:
+    def __init__(self, node: BaseNode):
+        self._node = node
+        self._huggingface_repo_parameter = HuggingFaceRepoParameter(
+            node,
+            repo_ids=[
+                "black-forest-labs/FLUX.1-Kontext-dev",
+            ],
+        )
+        self._seed_parameter = SeedParameter(node)
+
+    def add_input_parameters(self) -> None:
+        self._huggingface_repo_parameter.add_input_parameters()
+        self._node.add_parameter(
+            Parameter(
+                name="text_encoder",
+                input_types=["str"],
+                type="str",
+                allowed_modes=set(),
+                tooltip="text_encoder",
+                default_value="openai/clip-vit-large-patch14",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="text_encoder_2",
+                input_types=["str"],
+                type="str",
+                allowed_modes=set(),
+                tooltip="text_encoder_2",
+                default_value="google/t5-v1_1-xxl",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="image",
+                input_types=["ImageArtifact", "ImageUrlArtifact"],
+                type="ImageArtifact",
+                tooltip="Input image for img2img generation (optional)",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="prompt",
+                default_value="",
+                input_types=["str"],
+                type="str",
+                tooltip="prompt",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="prompt_2",
+                input_types=["str"],
+                type="str",
+                tooltip="optional prompt_2 - defaults to prompt",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="negative_prompt",
+                default_value="",
+                input_types=["str"],
+                type="str",
+                tooltip="optional negative_prompt",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="negative_prompt_2",
+                input_types=["str"],
+                type="str",
+                tooltip="optional negative_prompt_2 - defaults to negative_prompt",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="true_cfg_scale",
+                default_value=1.0,
+                input_types=["float"],
+                type="float",
+                tooltip="true_cfg_scale",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="width_by_height",
+                default_value="1024x1024",
+                input_types=["str"],
+                type="str",
+                tooltip="widthxheight",
+                traits={
+                    Options(
+                        choices=[
+                            "688x1504",
+                            "720x1456",
+                            "752x1392",
+                            "800x1328",
+                            "832x1248",
+                            "880x1184",
+                            "944x1104",
+                            "1024x1024",
+                            "1104x944",
+                            "1184x880",
+                            "1248x832",
+                            "1328x800",
+                            "1392x752",
+                            "1456x720",
+                            "1504x688",
+                        ]
+                    )
+                },
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="num_inference_steps",
+                default_value=28,
+                input_types=["int"],
+                type="int",
+                tooltip="num_inference_steps",
+            )
+        )
+        self._node.add_parameter(
+            Parameter(
+                name="guidance_scale",
+                default_value=3.5,
+                input_types=["float"],
+                type="float",
+                tooltip="guidance_scale",
+            )
+        )
+        self._seed_parameter.add_input_parameters()
+
+    def add_output_parameters(self) -> None:
+        self._node.add_parameter(
+            Parameter(
+                name="output_image",
+                output_type="ImageArtifact",
+                tooltip="The output image",
+                allowed_modes={ParameterMode.OUTPUT},
+            )
+        )
+
+    def validate_before_node_run(self) -> list[Exception] | None:
+        errors = self._huggingface_repo_parameter.validate_before_node_run() or []
+        # Check for if prompt exists:
+        prompt_exists = self.get_prompt() or self.get_prompt_2()
+        negative_prompt_exists = self.get_negative_prompt() or self.get_negative_prompt_2()
+        image_exists = self.get_input_image_pil()
+        if not prompt_exists and not negative_prompt_exists and not image_exists:
+            errors.append(ValueError("At least one prompt, negative prompt, or image must be provided."))
+
+        # Validate dimensions based on diffusers source logic
+        width = self.get_width()
+        height = self.get_height()
+
+        if width * height > (1024 * 1024):
+            errors.append(ValueError(f"Width ({width}) * Height ({height}) must be less than 1024*1024"))
+
+        # Get actual VAE scale factor from model config
+        try:
+            repo_id, revision = self.get_repo_revision()
+            vae_config = diffusers.AutoencoderKLConfig.from_pretrained(
+                pretrained_model_name_or_path=repo_id, revision=revision, local_files_only=True, subfolder="vae"
+            )
+            vae_scale_factor = 2 ** (len(vae_config.block_out_channels) - 1)
+            # Flux latents are packed into 2x2 patches, so multiply by 2
+            multiple_of = vae_scale_factor * 2
+        except Exception:
+            # Fallback to standard Flux values if model loading fails
+            multiple_of = 16
+
+        if width % multiple_of != 0:
+            errors.append(ValueError(f"Width ({width}) must be divisible by {multiple_of}"))
+
+        if height % multiple_of != 0:
+            errors.append(ValueError(f"Height ({height}) must be divisible by {multiple_of}"))
+
+        if width <= 0:
+            errors.append(ValueError(f"Width must be positive, got {width}"))
+
+        if height <= 0:
+            errors.append(ValueError(f"Height must be positive, got {height}"))
+
+        return errors or None
+
+    def after_value_set(self, parameter: Parameter, value: Any) -> None:
+        self._seed_parameter.after_value_set(parameter, value)
+
+    def preprocess(self) -> None:
+        self._seed_parameter.preprocess()
+
+    def get_repo_revision(self) -> tuple[str, str]:
+        return self._huggingface_repo_parameter.get_repo_revision()
+
+    def get_input_image_pil(self) -> Image | None:
+        image_artifact = self._node.get_parameter_value("image")
+        if image_artifact is None:
+            return None
+        if isinstance(image_artifact, ImageUrlArtifact):
+            image_artifact = load_image_from_url_artifact(image_artifact)
+        input_image_pil = image_artifact_to_pil(image_artifact)
+        return input_image_pil.convert("RGB")
+
+    def get_prompt(self) -> str:
+        return self._node.get_parameter_value("prompt")
+
+    def get_prompt_2(self) -> str:
+        prompt_2 = self._node.get_parameter_value("prompt_2")
+        if prompt_2 is None:
+            return self.get_prompt()
+        return prompt_2
+
+    def get_negative_prompt(self) -> str:
+        return self._node.get_parameter_value("negative_prompt")
+
+    def get_negative_prompt_2(self) -> str:
+        negative_prompt_2 = self._node.get_parameter_value("negative_prompt_2")
+        if negative_prompt_2 is None:
+            return self.get_negative_prompt()
+        return negative_prompt_2
+
+    def get_true_cfg_scale(self) -> float:
+        return float(self._node.get_parameter_value("true_cfg_scale"))
+
+    def get_width(self) -> int:
+        width_by_height = self._node.get_parameter_value("width_by_height")
+        return int(width_by_height.split("x")[0])
+
+    def get_height(self) -> int:
+        width_by_height = self._node.get_parameter_value("width_by_height")
+        return int(width_by_height.split("x")[1])
+
+    def get_num_inference_steps(self) -> int:
+        return int(self._node.get_parameter_value("num_inference_steps"))
+
+    def get_guidance_scale(self) -> float:
+        return float(self._node.get_parameter_value("guidance_scale"))
+
+    def get_effective_size(self) -> tuple[int, int]:
+        """Get the effective width and height."""
+        return self.get_width(), self.get_height()
+
+    def get_pipe_kwargs(self) -> dict:
+        return {
+            "image": self.get_input_image_pil(),
+            "prompt": self.get_prompt(),
+            "prompt_2": self.get_prompt_2(),
+            "negative_prompt": self.get_negative_prompt(),
+            "negative_prompt_2": self.get_negative_prompt_2(),
+            "true_cfg_scale": self.get_true_cfg_scale(),
+            "width": self.get_width(),
+            "height": self.get_height(),
+            "num_inference_steps": self.get_num_inference_steps(),
+            "guidance_scale": self.get_guidance_scale(),
+            "generator": self._seed_parameter.get_generator(),
+        }
+
+    def publish_output_image_preview_placeholder(self) -> None:
+        width, height = self.get_effective_size()
+        preview_placeholder_image = PIL.Image.new("RGB", (width, height), color="black")
+        self._node.publish_update_to_parameter("output_image", pil_to_image_artifact(preview_placeholder_image))
+
+    def latents_to_image_pil(self, pipe: diffusers.FluxKontextPipeline, latents: Any) -> Image:
+        width, height = self.get_effective_size()
+        latents = pipe._unpack_latents(latents, height, width, pipe.vae_scale_factor)
+        latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
+        image = pipe.vae.decode(latents, return_dict=False)[0]
+        intermediate_pil_image = pipe.image_processor.postprocess(image, output_type="pil")[0]
+        return intermediate_pil_image
+
+    def publish_output_image_preview_latents(self, pipe: diffusers.FluxKontextPipeline, latents: Any) -> None:
+        preview_image_pil = self.latents_to_image_pil(pipe, latents)
+        preview_image_artifact = pil_to_image_artifact(preview_image_pil)
+        self._node.publish_update_to_parameter("output_image", preview_image_artifact)
+
+    def publish_output_image(self, output_image_pil: Image) -> None:
+        image_artifact = pil_to_image_artifact(output_image_pil)
+        self._node.set_parameter_value("output_image", image_artifact)
+        self._node.parameter_output_values["output_image"] = image_artifact