fixing pre-existing make pipeline issues

modular_diffusion_nodes_library/artifact_utils/latent_artifact.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import pickle
-from typing import Any
+from typing import Any, NoReturn
 
 import torch  # type: ignore[reportMissingImports]
 from griptape.artifacts.base_artifact import BaseArtifact
@@ -106,7 +106,7 @@ def to_torch(
     def to_text(self) -> str:
         return repr(self)
 
-    def to_dict(self) -> dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:  # type: ignore[reportIncompatibleMethodOverride]
         tensor = self.to_torch().detach().cpu()
         if tensor.numel() == 0:
             return {
@@ -137,7 +137,7 @@ def to_dict(self) -> dict[str, Any]:
             "sample": sample,
         }
 
-    def __reduce__(self) -> None:
+    def __reduce__(self) -> NoReturn:
         raise pickle.PicklingError("LatentArtifact is an in-process-only type and should not be serialized.")
 
     def __bool__(self) -> bool:

modular_diffusion_nodes_library/artifact_utils/pipeline_artifact.py

@@ -107,6 +107,7 @@ def __init__(
         self._is_prequantized = is_prequantized
         self._supports_layerwise_casting = supports_layerwise_casting
         self._requires_device_map = requires_device_map
+        self._extra_runtime_adapter_steps: list[PipelineRuntimeAdapterStep] = []
 
     @property
     def builder_module(self) -> str | None:
@@ -174,7 +175,7 @@ def get_or_build_pipeline(self, log_params: Any | None = None) -> ModularPipelin
 
     def runtime_adapter_steps(self) -> list[PipelineRuntimeAdapterStep]:
         """Per-generation context-managed transformations applied around each generation call."""
-        return list(self._extra_runtime_adapter_steps) if hasattr(self, "_extra_runtime_adapter_steps") else []
+        return list(self._extra_runtime_adapter_steps)
 
     def with_additional_runtime_adapter_steps(
         self, steps: list[PipelineRuntimeAdapterStep]

modular_diffusion_nodes_library/artifact_utils/pipeline_build_steps.py

@@ -147,7 +147,7 @@ def apply(self, pipe: Pipe | None, *, log_params: Any | None = None) -> Pipe:
             raise RuntimeError("ApplyOptimizationStep requires a pipe.")
         with _profile(log_params, "Applying optimizations"):
             optimize_diffusion_pipeline(
-                pipe=pipe,
+                pipe=pipe,  # type: ignore[reportArgumentType]
                 is_prequantized=self._is_prequantized,
                 supports_layerwise_casting=self._supports_layerwise_casting,
                 requires_device_map=self._requires_device_map,

modular_diffusion_nodes_library/latent_pipeline_drivers/base_driver.py

@@ -149,7 +149,7 @@ def _call_block(self, block: ModularPipelineBlocks, **kwargs: Any) -> dict[str,
         for param in block.inputs:
             if param.name in kwargs:
                 state.set(param.name, kwargs[param.name], param.kwargs_type)
-        _, state = block(self.modular_pipe, state)
+        _, state = block(self.modular_pipe, state)  # type: ignore[reportOperatorIssue]
         return state.values
 
     @staticmethod
@@ -261,7 +261,13 @@ def encode_masked_image(
     ) -> LatentArtifact:
         """Encode the source image with the masked region zeroed out."""
         image_processor = self.modular_pipe.image_processor
-        source_t = image_processor.preprocess(image.image, height=image.image.height, width=image.image.width)
+        if isinstance(image.image, torch.Tensor):
+            height = image.image.shape[-2]
+            width = image.image.shape[-1]
+        else:
+            height = image.image.height
+            width = image.image.width
+        source_t = image_processor.preprocess(image.image, height=height, width=width)
         mask_t = torch.from_numpy(np.array(mask.mask, dtype="float32") / 255.0)[None, None]
         masked_t = source_t * (mask_t < 0.5)
         return self.encode_media(ImageMedia(image=masked_t, source_shape=image.source_shape), generator_state)
@@ -436,10 +442,11 @@ def _get_inpaint_kwargs(self, artifact: InpaintMaskArtifact) -> dict[str, Any]:
         """Map an InpaintMaskArtifact to pipeline call kwargs."""
         device, dtype = self._get_device_and_type()
         result: dict[str, Any] = {
-            "image": artifact.source_latent.to(device=device, dtype=dtype),
             "mask_image": artifact.mask_image,
             "strength": artifact.strength,
         }
+        if artifact.source_latent is not None:
+            result["image"] = artifact.source_latent.to(device=device, dtype=dtype)
         if artifact.masked_latent is not None:
             result["masked_image_latents"] = artifact.masked_latent.to(device=device, dtype=dtype)
         return result

modular_diffusion_nodes_library/latent_pipeline_drivers/flux.py

@@ -2,13 +2,8 @@
 from typing import Any, ClassVar, override
 
 import torch  # type: ignore[reportMissingImports]
-from diffusers import (  # type: ignore[reportMissingImports]
-    FluxControlNetInpaintPipeline,
-    FluxControlNetModel,
-    FluxControlNetPipeline,
-    FluxInpaintPipeline,
-)
 from diffusers.models import FluxTransformer2DModel  # type: ignore[reportMissingImports]
+from diffusers.models.controlnets.controlnet_flux import FluxControlNetModel  # type: ignore[reportMissingImports]
 from diffusers.modular_pipelines.flux.before_denoise import (  # type: ignore[reportMissingImports]
     FluxImg2ImgPrepareLatentsStep,
     FluxImg2ImgSetTimestepsStep,
@@ -25,6 +20,13 @@
 )
 from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec  # type: ignore[reportMissingImports]
 from diffusers.pipelines.flux.pipeline_flux import FluxPipeline  # type: ignore[reportMissingImports]
+from diffusers.pipelines.flux.pipeline_flux_controlnet import (
+    FluxControlNetPipeline,  # type: ignore[reportMissingImports]
+)
+from diffusers.pipelines.flux.pipeline_flux_controlnet_inpainting import (
+    FluxControlNetInpaintPipeline,  # type: ignore[reportMissingImports]
+)
+from diffusers.pipelines.flux.pipeline_flux_inpaint import FluxInpaintPipeline  # type: ignore[reportMissingImports]
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline  # type: ignore[reportMissingImports]
 from PIL.Image import Image
 
@@ -165,7 +167,7 @@ def create_noise_latent(self, source_shape: tuple[int, ...], generator_state: Ge
             batch_size=1,
             dtype=dtype,
         )
-        latents = output_state.get("latents")
+        latents = self._get_required(output_state, "latents", torch.Tensor)
         latents = latents.to(device)
         latents = self.unpack_latents(latents, height, width)
         return self._make_latent_artifact(
@@ -213,7 +215,7 @@ def add_noise_to_latent(
             image_latents=packed_image_latents,
         )
 
-        noisy_latents = output_state.get("latents")
+        noisy_latents = self._get_required(output_state, "latents", torch.Tensor)
         unpacked = self.unpack_latents(noisy_latents, height, width)
         return self._make_latent_artifact(
             unpacked,
@@ -233,7 +235,8 @@ def decode_latent(self, latent: LatentArtifact) -> Image:
 
         decode_pipeline = self.modular_pipe.blocks.sub_blocks["decode"]
         output_state = self._call_block(decode_pipeline, latents=packed, output_type="pil", width=width, height=height)
-        return output_state.get("images")[0]
+        images = self._get_required(output_state, "images", list)
+        return images[0]
 
     @override
     def encode_media(self, media: ImageMedia | VideoMedia, generator_state: GeneratorState) -> LatentArtifact:
@@ -249,7 +252,7 @@ def encode_media(self, media: ImageMedia | VideoMedia, generator_state: Generato
         generator = generator_state.to_generator()
         encode_pipeline = self.modular_pipe.blocks.sub_blocks["vae_encoder"]
         output_state = self._call_block(encode_pipeline, image=image, height=height, width=width, generator=generator)
-        latents = output_state.get("image_latents")
+        latents = self._get_required(output_state, "image_latents", torch.Tensor)
         return self._make_latent_artifact(latents, source_shape=media.source_shape)
 
     @override

modular_diffusion_nodes_library/latent_pipeline_drivers/flux2_base.py

@@ -62,7 +62,7 @@ def create_noise_latent(self, source_shape: tuple[int, ...], generator_state: Ge
 
         unpack_latents = Flux2UnpackLatentsStep()
         latents_state = self._call_block(unpack_latents, latents=packed_latents, latent_ids=latent_ids)
-        latents = latents_state.get("latents")
+        latents = self._get_required(latents_state, "latents", torch.Tensor)
         latents = latents.to(device)
         return self._make_latent_artifact(
             latents,
@@ -78,7 +78,7 @@ def decode_latent(self, latent: LatentArtifact) -> Image:
         decode_block = self.modular_pipe.blocks.sub_blocks["decode"]
         output_state = self._call_block(decode_block, latents=latents, output_type="pil")
 
-        return output_state.get("images")[0]
+        return self._get_required(output_state, "images", list)[0]
 
     def _unpack_latents(self, latents: torch.Tensor, height: int, width: int) -> torch.Tensor:
         """Convert packed Flux2 latents ``[B, seq_len, C]`` to ``[B, C, H/vae/2, W/vae/2]``."""
@@ -97,7 +97,7 @@ def _unpack_latents(self, latents: torch.Tensor, height: int, width: int) -> tor
         latent_ids = id_state.get("latent_ids")
 
         unpack_state = self._call_block(Flux2UnpackLatentsStep(), latents=latents, latent_ids=latent_ids)
-        return unpack_state.get("latents")
+        return self._get_required(unpack_state, "latents", torch.Tensor)
 
     @override
     def encode_media(self, media: ImageMedia | VideoMedia, generator_state: GeneratorState) -> LatentArtifact:
@@ -113,7 +113,9 @@ def encode_media(self, media: ImageMedia | VideoMedia, generator_state: Generato
         output_state = self._call_block(
             encode_block, image=image, height=image.height, width=image.width, generator=generator
         )
-        return self._make_latent_artifact(output_state.get("image_latents")[0], source_shape=media.source_shape)
+        return self._make_latent_artifact(
+            self._get_required(output_state, "image_latents", list)[0], source_shape=media.source_shape
+        )
 
     @override
     def encode_masked_image(
@@ -126,17 +128,19 @@ def encode_masked_image(
 
         # PIL -> aligned, normalized tensor (multiple-of-32 alignment + <=1024^2 cap).
         pre = self._call_block(Flux2ProcessImagesInputStep(), image=pil_image)
-        source_t = pre.get("condition_images")[0]
+        source_t = self._get_required(pre, "condition_images", list)[0]
 
         # Resize mask to match the (possibly aligned/cropped) tensor dims.
         aligned_h, aligned_w = source_t.shape[-2:]
-        mask_resized = mask.mask.convert("L").resize((aligned_w, aligned_h), PIL.Image.NEAREST)
+        mask_resized = mask.mask.convert("L").resize((aligned_w, aligned_h), PIL.Image.Resampling.NEAREST)
         mask_t = TF.to_tensor(mask_resized)[None]  # (1, 1, H, W), float32 in [0, 1]
 
         masked_t = source_t * (mask_t < 0.5)
 
         out = self._call_block(Flux2VaeEncoderStep(), condition_images=[masked_t])
-        return self._make_latent_artifact(out.get("image_latents")[0], source_shape=image.source_shape)
+        return self._make_latent_artifact(
+            self._get_required(out, "image_latents", list)[0], source_shape=image.source_shape
+        )
 
     @override
     def add_noise_to_latent(

modular_diffusion_nodes_library/latent_pipeline_drivers/flux2_klein.py

@@ -55,7 +55,7 @@ def _get_inpaint_kwargs(self, artifact: InpaintMaskArtifact) -> dict[str, Any]:
         The pipeline will BN-normalize image latent internally, so we must denormalize first.
         """
         device, dtype = self._get_device_and_type()
-        source_latent = self._denormalize_latent(artifact.source_latent).to(device=device, dtype=dtype)
+        source_latent = self._denormalize_latent(artifact.source_latent).to(device=device, dtype=dtype)  # type: ignore[reportArgumentType]
         return {
             "image": source_latent,
             "mask_image": artifact.mask_image,

modular_diffusion_nodes_library/latent_pipeline_drivers/ltx.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, ClassVar, override
+from typing import Any, ClassVar, cast, override
 
 import torch  # type: ignore[reportMissingImports]
 from diffusers import LTXConditionPipeline  # type: ignore[reportMissingImports]
@@ -13,14 +13,16 @@
     LTXAutoVaeEncoderStep,
 )
 from diffusers.modular_pipelines.modular_pipeline import (  # type: ignore[reportMissingImports]
-    ComponentSpec,
-    InputParam,
     ModularPipeline,
     ModularPipelineBlocks,
-    OutputParam,
     PipelineState,
     SequentialPipelineBlocks,
 )
+from diffusers.modular_pipelines.modular_pipeline_utils import (  # type: ignore[reportMissingImports]
+    ComponentSpec,
+    InputParam,
+    OutputParam,
+)
 from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition  # type: ignore[reportMissingImports]
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline  # type: ignore[reportMissingImports]
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler  # type: ignore[reportMissingImports]
@@ -86,7 +88,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @torch.no_grad()
     def __call__(self, components, state: PipelineState):
-        block_state = self.get_block_state(state)
+        block_state = cast(Any, self.get_block_state(state))
 
         init_timestep = min(
             block_state.num_inference_steps * block_state.strength,
@@ -131,7 +133,7 @@ def intermediate_outputs(self) -> list[OutputParam]:
 
     @torch.no_grad()
     def __call__(self, components, state: PipelineState):
-        block_state = self.get_block_state(state)
+        block_state = cast(Any, self.get_block_state(state))
 
         latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
         block_state.latents = components.scheduler.scale_noise(
@@ -150,7 +152,7 @@ class LTXAddNoiseStep(SequentialPipelineBlocks):
     then strength slicing, then scale_noise.
     """
 
-    model_name = "ltx"
+    model_name = "ltx"  # type: ignore[reportIncompatibleMethodOverride]
     block_classes = [
         LTXSetTimestepsStep,
         LTXSetTimestepsWithStrengthStep,
@@ -236,7 +238,7 @@ def create_noise_latent(self, source_shape: tuple[int, ...], generator_state: Ge
             num_videos_per_prompt=1,
             generator=generator,
         )
-        packed_latents = output_state.get("latents")
+        packed_latents = self._get_required(output_state, "latents", torch.Tensor)
         latents = self._unpack_latents(packed_latents, height, width, num_frames)
         return self._make_latent_artifact(
             latents,
@@ -278,7 +280,7 @@ def add_noise_to_latent(
         )
         noise_generator_state = GeneratorState.from_artifact(noise_artifact) or generator_state
         return self._make_latent_artifact(
-            output_state.get("latents"),
+            self._get_required(output_state, "latents", torch.Tensor),
             source_shape=source_shape,
             upstream=latent,
             meta=noise_generator_state.as_meta(),
@@ -312,16 +314,17 @@ def decode_latent(self, latent: LatentArtifact) -> DecodeResult:
             output_type="pil",
             decode_timestep=0.0,
         )
-        videos = output_state.get("videos")
-        return videos[0]
+        return self._get_required(output_state, "videos", list)[0]
 
     @override
     def encode_media(self, media: ImageMedia | VideoMedia, generator_state: GeneratorState) -> LatentArtifact:
         """Encode an image or video into a 5D latent ``[B, C, T, H, W]``."""
         generator = generator_state.to_generator()
         if isinstance(media, ImageMedia):
             output_state = self._call_block(LTXAutoVaeEncoderStep(), image=media.image, generator=generator)
-            return self._make_latent_artifact(output_state.get("image_latents"), source_shape=media.source_shape)
+            return self._make_latent_artifact(
+                self._get_required(output_state, "image_latents", torch.Tensor), source_shape=media.source_shape
+            )
 
         device, dtype = self._get_device_and_type()
         vae = self.modular_pipe.vae

modular_diffusion_nodes_library/latent_pipeline_drivers/ltx2.py

@@ -283,7 +283,7 @@ def _extract_latents_from_output(self, pipe_output: Any) -> torch.Tensor:
         """LTX2 pipelines return video frames under ``.frames`` instead of ``.images``."""
         return pipe_output.frames
 
-    def _decode_hdr_to_linear_np(self, video: torch.Tensor) -> np.ndarray:
+    def _decode_hdr_to_linear_np(self, video: torch.Tensor) -> torch.Tensor | np.ndarray:
         """Convert decoded LogC3-compressed VAE output to linear HDR np.ndarray.
 
         Returns shape ``(B, F, H, W, 3)`` float32 with linear HDR values in ``[0, ∞)``.
@@ -328,7 +328,7 @@ def decode_latent(self, latent: LatentArtifact) -> list[Image] | np.ndarray:
 
         if self._latent_was_produced_for_hdr(latent):
             # HDR IC-LoRA path: return raw linear HDR for encode_hdr_tensor_to_mp4 in vae_decoder.
-            return self._decode_hdr_to_linear_np(video)
+            return self._decode_hdr_to_linear_np(video)  # type: ignore[reportReturnType]
 
         frames = self.pipe.video_processor.postprocess_video(video, output_type="pil")[0]
         return frames