Limitex-ComfyUI-Diffusers/src/infra/diffusers/_stream_diffusion_repository.py

import copy
from typing import Any

import torch
from comfy.model_management import get_torch_device  # pyright: ignore[reportMissingImports]
from diffusers import AutoencoderTiny, StableDiffusionPipeline
from PIL import Image
from streamdiffusion import StreamDiffusion  # type: ignore[import-untyped]
from streamdiffusion.image_utils import postprocess_image  # type: ignore[import-untyped]

from ...domain.model import (
    CFGScale,
    CFGType,
    Delta,
    FrameBufferSize,
    ImageSize,
    NumSamples,
    Seed,
    Steps,
    TIndexList,
    WarmupCount,
)
from ...domain.repositories import StreamDiffusionRepository
from ._cache import get_cache_path

LINEAR_NDIM = 2
CONV_NDIM_THRESHOLD = 3


class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
    def __init__(self) -> None:
        self.device = get_torch_device()
        self.cache_dir = get_cache_path()
        self.dtype = torch.float16

    def create_stream(
        self,
        pipeline: StableDiffusionPipeline,
        t_index_list: TIndexList,
        image_size: ImageSize,
        do_add_noise: bool,
        use_denoising_batch: bool,
        frame_buffer_size: FrameBufferSize,
        cfg_type: CFGType,
        lcm_lora_weights: dict[str, Any],
        tiny_vae_name: str,
        enable_xformers: bool,
    ) -> Any:
        """Create a Stream Diffusion stream instance."""
        # Deep copy to avoid modifying the original pipeline
        # Note: load_lcm_lora() and fuse_lora() modify the pipeline,
        # so we need to copy it to allow reusing the original pipeline in ComfyUI
        pipeline_copy = copy.deepcopy(pipeline)
        lora_weights_copy = self._filter_and_reshape_lora_weights(
            copy.deepcopy(lcm_lora_weights), pipeline_copy
        )

        # Create stream
        stream = StreamDiffusion(
            pipe=pipeline_copy,
            t_index_list=t_index_list.indices,
            torch_dtype=self.dtype,
            width=image_size.width,
            height=image_size.height,
            do_add_noise=do_add_noise,
            use_denoising_batch=use_denoising_batch,
            frame_buffer_size=frame_buffer_size.value,
            cfg_type=cfg_type.value.value,  # Get the string value from enum
        )

        # Load and fuse LCM LoRA
        # Pass ignore_mismatched_sizes to tolerate conv weight shape differences (e.g., meta tensors)
        stream.load_lcm_lora(
            lora_weights_copy,
            low_cpu_mem_usage=False,
            ignore_mismatched_sizes=True,
        )
        stream.fuse_lora()

        # Load tiny VAE
        stream.vae = AutoencoderTiny.from_pretrained(  # type: ignore[no-untyped-call]
            pretrained_model_name_or_path=tiny_vae_name,
            torch_dtype=self.dtype,
            cache_dir=self.cache_dir,
        ).to(device=pipeline_copy.device, dtype=pipeline_copy.dtype)  # type: ignore[attr-defined]

        # Enable xformers if requested
        if enable_xformers:
            pipeline_copy.enable_xformers_memory_efficient_attention()  # type: ignore[attr-defined]

        return stream

    def warmup_stream(
        self,
        stream: Any,
        negative_prompt: str,
        steps: Steps,
        cfg: CFGScale,
        delta: Delta,
        seed: Seed,
        warmup_count: WarmupCount,
    ) -> None:
        """Warm up the stream with given parameters."""
        stream.prepare(
            prompt="",
            negative_prompt=negative_prompt,
            num_inference_steps=steps.value,
            guidance_scale=cfg.value,
            delta=delta.value,
            seed=seed.value,
        )

        for _ in range(warmup_count.value):
            stream()

    def update_prompt(self, stream: Any, prompt: str) -> None:
        """Update the prompt for the stream."""
        stream.update_prompt(prompt)

    def sample_txt2img(
        self,
        stream: Any,
        num_samples: NumSamples,
    ) -> list[Image.Image]:
        """Generate images using txt2img."""
        result: list[Image.Image] = []
        for _ in range(num_samples.value):
            x_output = stream.txt2img()
            image = postprocess_image(x_output, output_type="pil")[0]
            result.append(image)
        return result

    def sample_with_images(
        self,
        stream: Any,
        prompt: str,
        negative_prompt: str,
        steps: Steps,
        cfg: CFGScale,
        delta: Delta,
        seed: Seed,
        num_samples: NumSamples,
        warmup_count: WarmupCount,
        input_images: list[Image.Image] | None,
    ) -> list[Image.Image]:
        """Generate images with optional input images."""
        # Prepare stream
        stream.prepare(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps.value,
            guidance_scale=cfg.value,
            delta=delta.value,
            seed=seed.value,
        )

        # Resize input images if provided
        if input_images is not None:
            resized_images = [img.resize((stream.width, stream.height)) for img in input_images]
        else:
            resized_images = None

        # Warmup
        for _ in range(warmup_count.value):
            stream()

        # Generate images
        result: list[Image.Image] = []
        if resized_images is None:
            # Text-to-image: Generate num_samples images
            for _ in range(num_samples.value):
                x_output = stream.txt2img()
                image = postprocess_image(x_output, output_type="pil")[0]
                result.append(image)
        else:
            # Image-to-image: Process each input image num_samples times
            for _ in range(num_samples.value):
                for img in resized_images:
                    x_output = stream(img)
                    image = postprocess_image(x_output, output_type="pil")[0]
                    result.append(image)

        return result

    @staticmethod
    def _filter_and_reshape_lora_weights(
        lora_weights: dict[str, Any], pipeline: StableDiffusionPipeline
    ) -> dict[str, Any]:
        """Keep only compatible LoRA tensors and reshape 2D conv-proj weights as needed."""
        reshaped: dict[str, Any] = {}
        # SD1.x UNet channel sizes + text encoder hidden size inform allowed dims
        allowed_in_dims: set[int] = set(pipeline.unet.config.block_out_channels)  # type: ignore[attr-defined]
        if hasattr(pipeline.text_encoder.config, "hidden_size"):  # type: ignore[attr-defined]
            allowed_in_dims.add(pipeline.text_encoder.config.hidden_size)  # type: ignore[attr-defined]
        # SDXL uses dual text encoders; include second encoder dims when present
        if hasattr(pipeline, "text_encoder_2") and hasattr(
            pipeline.text_encoder_2.config, "hidden_size"
        ):
            allowed_in_dims.add(pipeline.text_encoder_2.config.hidden_size)
        # Allow doubled dims for concat projections (e.g., 2560 = 2 * 1280)
        allowed_conv_in_dims: set[int] = allowed_in_dims.union({c * 2 for c in allowed_in_dims})
        # Only attention projection convs in UNet (proj_in/proj_out) need 1x1 conv shape.
        conv_like_markers = ("proj_in", "proj_out")

        # First, copy non-LoRA keys verbatim
        for key, value in lora_weights.items():
            # Drop alpha scaling entries to avoid conversion failures
            if key.endswith(".alpha"):
                continue
            if "lora_down.weight" not in key and "lora_up.weight" not in key:
                reshaped[key] = value

        # Then process LoRA pairs, ensuring both down/up are present and compatible
        for key, value in lora_weights.items():
            if not key.endswith("lora_down.weight"):
                continue
            up_key = key.replace("lora_down.weight", "lora_up.weight")
            down_tensor = value
            up_tensor = lora_weights.get(up_key)

            # Require matching pair
            if not isinstance(down_tensor, torch.Tensor) or not isinstance(up_tensor, torch.Tensor):
                continue

            # Drop obviously incompatible LoRA tensors (e.g., SDXL-sized on SD1.5 pipeline)
            # 2D (linear) uses allowed_in_dims; conv uses allowed_conv_in_dims.
            if down_tensor.ndim == LINEAR_NDIM and (down_tensor.shape[1] not in allowed_in_dims):
                continue
            if up_tensor.ndim == LINEAR_NDIM and (up_tensor.shape[1] not in allowed_in_dims):
                continue
            if down_tensor.ndim >= CONV_NDIM_THRESHOLD and (
                down_tensor.shape[1] not in allowed_conv_in_dims
            ):
                continue
            if up_tensor.ndim >= CONV_NDIM_THRESHOLD and (
                up_tensor.shape[1] not in allowed_conv_in_dims
            ):
                continue

            # Reshape conv-projection weights; keep others as-is
            if down_tensor.ndim == LINEAR_NDIM and any(
                marker in key for marker in conv_like_markers
            ):
                down_tensor = down_tensor.unsqueeze(-1).unsqueeze(-1)
            if up_tensor.ndim == LINEAR_NDIM and any(
                marker in up_key for marker in conv_like_markers
            ):
                up_tensor = up_tensor.unsqueeze(-1).unsqueeze(-1)

            reshaped[key] = down_tensor
            reshaped[up_key] = up_tensor

        return reshaped