comfy-diffusion 0.1.0__py3-none-any.whl

comfy_diffusion/__init__.py

@@ -0,0 +1,30 @@
+"""Public package entrypoint for comfy_diffusion."""
+
+from ._runtime import ensure_comfyui_on_path
+from .lora import apply_lora
+from .runtime import check_runtime
+from .vae import (
+    vae_decode,
+    vae_decode_batch,
+    vae_decode_batch_tiled,
+    vae_decode_tiled,
+    vae_encode,
+    vae_encode_batch,
+    vae_encode_batch_tiled,
+    vae_encode_tiled,
+)
+
+ensure_comfyui_on_path()
+
+__all__ = [
+    "check_runtime",
+    "vae_decode",
+    "vae_decode_batch",
+    "vae_decode_batch_tiled",
+    "vae_decode_tiled",
+    "vae_encode",
+    "vae_encode_batch",
+    "vae_encode_batch_tiled",
+    "vae_encode_tiled",
+    "apply_lora",
+]

comfy_diffusion/_runtime.py

@@ -0,0 +1,26 @@
+"""Internal runtime bootstrap for comfy_diffusion.
+
+Path insertion is intentionally lightweight and import-safe: this module must not
+import torch or comfy internals just to make ComfyUI discoverable.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+def _comfyui_root() -> Path:
+    """Return the absolute path to the vendored ComfyUI directory."""
+    return Path(__file__).resolve().parents[1] / "vendor" / "ComfyUI"
+
+
+def ensure_comfyui_on_path() -> Path:
+    """Ensure vendored ComfyUI is importable and return the inserted path."""
+    comfyui_root = _comfyui_root()
+    comfyui_root_str = str(comfyui_root)
+
+    if comfyui_root_str not in sys.path:
+        sys.path.insert(0, comfyui_root_str)
+
+    return comfyui_root

comfy_diffusion/audio.py

@@ -0,0 +1,168 @@
+"""Audio helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, cast
+
+
+class _LtxvAudioVaeEncoder(Protocol):
+    sample_rate: int
+
+    def encode(self, audio: Any) -> Any: ...
+
+
+class _LtxvAudioVaeDecoder(Protocol):
+    output_sample_rate: int
+
+    def decode(self, latent: Any) -> Any: ...
+
+
+class _LtxvAudioVae(Protocol):
+    sample_rate: int
+    latent_channels: int
+    latent_frequency_bins: int
+
+    def num_of_latents_from_frames(self, frames_number: int, frame_rate: int) -> int: ...
+
+
+class _AceStep15Clip(Protocol):
+    def tokenize(
+        self,
+        tags: str,
+        *,
+        lyrics: str,
+        bpm: int,
+        duration: float,
+        timesignature: int,
+        language: str,
+        keyscale: str,
+        seed: int,
+        generate_audio_codes: bool,
+        cfg_scale: float,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        min_p: float,
+    ) -> Any: ...
+
+    def encode_from_tokens_scheduled(self, tokens: Any) -> Any: ...
+
+
+def _get_ltxv_empty_latent_audio_type() -> Any:
+    """Resolve ComfyUI LTXVEmptyLatentAudio node at call time."""
+    from ._runtime import ensure_comfyui_on_path
+
+    ensure_comfyui_on_path()
+    from comfy_extras.nodes_lt_audio import LTXVEmptyLatentAudio
+
+    return LTXVEmptyLatentAudio
+
+
+def _get_ace_step_15_latent_audio_dependencies() -> tuple[Any, Any]:
+    """Resolve torch and ComfyUI model management at call time."""
+    from ._runtime import ensure_comfyui_on_path
+
+    ensure_comfyui_on_path()
+    import comfy.model_management
+    import torch
+
+    return torch, comfy.model_management
+
+
+def _unwrap_node_output(output: Any) -> Any:
+    """Return first output for ComfyUI V3 nodes and tuple-style APIs."""
+    if hasattr(output, "result"):
+        return output.result[0]
+    if isinstance(output, tuple):
+        return output[0]
+    return output
+
+
+def ltxv_audio_vae_encode(vae: _LtxvAudioVaeEncoder, audio: Any) -> dict[str, Any]:
+    """Encode raw audio with an LTXV audio VAE."""
+    audio_latents = vae.encode(audio)
+    return {"samples": audio_latents, "sample_rate": int(vae.sample_rate), "type": "audio"}
+
+
+def ltxv_audio_vae_decode(vae: _LtxvAudioVaeDecoder, latent: Any) -> dict[str, Any]:
+    """Decode latent audio with an LTXV audio VAE."""
+    latent_tensor = latent["samples"] if isinstance(latent, dict) else latent
+    if getattr(latent_tensor, "is_nested", False):
+        latent_tensor = latent_tensor.unbind()[-1]
+    audio = vae.decode(latent_tensor).to(latent_tensor.device)
+    return {"waveform": audio, "sample_rate": int(vae.output_sample_rate)}
+
+
+def ltxv_empty_latent_audio(
+    audio_vae: _LtxvAudioVae,
+    frames_number: int,
+    frame_rate: int = 25,
+    batch_size: int = 1,
+) -> dict[str, Any]:
+    """Create empty LTXV audio latents compatible with ComfyUI's audio pipeline."""
+    ltxv_empty_latent_audio_type = _get_ltxv_empty_latent_audio_type()
+    return cast(
+        dict[str, Any],
+        _unwrap_node_output(
+            ltxv_empty_latent_audio_type.execute(
+                frames_number=frames_number,
+                frame_rate=frame_rate,
+                batch_size=batch_size,
+                audio_vae=audio_vae,
+            )
+        )
+    )
+
+
+def encode_ace_step_15_audio(
+    clip: _AceStep15Clip,
+    tags: str,
+    lyrics: str = "",
+    seed: int = 0,
+    bpm: int = 120,
+    duration: float = 120.0,
+    timesignature: str = "4",
+    language: str = "en",
+    keyscale: str = "C major",
+    generate_audio_codes: bool = True,
+    cfg_scale: float = 2.0,
+    temperature: float = 0.85,
+    top_p: float = 0.9,
+    top_k: int = 0,
+    min_p: float = 0.0,
+) -> Any:
+    """Encode ACE Step 1.5 text/audio metadata conditioning."""
+    tokens = clip.tokenize(
+        tags,
+        lyrics=lyrics,
+        bpm=bpm,
+        duration=duration,
+        timesignature=int(timesignature),
+        language=language,
+        keyscale=keyscale,
+        seed=seed,
+        generate_audio_codes=generate_audio_codes,
+        cfg_scale=cfg_scale,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        min_p=min_p,
+    )
+    return clip.encode_from_tokens_scheduled(tokens)
+
+
+def empty_ace_step_15_latent_audio(seconds: float, batch_size: int = 1) -> dict[str, Any]:
+    """Create empty ACE Step 1.5 latents used as sampler noise input."""
+    torch, model_management = _get_ace_step_15_latent_audio_dependencies()
+    length = round(seconds * 48000 / 1920)
+    latent = torch.zeros([batch_size, 64, length], device=model_management.intermediate_device())
+    return {"samples": latent, "type": "audio"}
+
+
+__all__ = [
+    "ltxv_audio_vae_encode",
+    "ltxv_audio_vae_decode",
+    "ltxv_empty_latent_audio",
+    "encode_ace_step_15_audio",
+    "empty_ace_step_15_latent_audio",
+]

comfy_diffusion/conditioning.py

@@ -0,0 +1,25 @@
+"""Prompt conditioning helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+
+class _ClipTextEncoder(Protocol):
+    def tokenize(self, text: str) -> Any: ...
+
+    def encode_from_tokens_scheduled(self, tokens: Any) -> Any: ...
+
+
+def encode_prompt(clip: _ClipTextEncoder, text: str) -> Any:
+    """Encode prompt text with a ComfyUI-compatible CLIP object.
+
+    Positive and negative prompts use the same encoding path; prompt
+    semantics are owned by the caller.
+    """
+    normalized_text = " " if text == "" else text
+    tokens = clip.tokenize(normalized_text)
+    return clip.encode_from_tokens_scheduled(tokens)
+
+
+__all__ = ["encode_prompt"]

comfy_diffusion/lora.py

@@ -0,0 +1,34 @@
+"""LoRA application helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, cast
+
+
+def apply_lora(
+    model: Any,
+    clip: Any,
+    path: str | Path,
+    strength_model: float,
+    strength_clip: float,
+) -> tuple[Any, Any]:
+    """Apply a LoRA file to a model/CLIP pair and return patched copies.
+
+    The returned pair can be passed back into ``apply_lora`` to stack
+    multiple LoRAs by chaining calls.
+    """
+    from ._runtime import ensure_comfyui_on_path
+
+    ensure_comfyui_on_path()
+
+    import comfy.sd
+    import comfy.utils
+
+    lora_path = str(Path(path))
+    lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
+    patched = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
+    return cast(tuple[Any, Any], patched)
+
+
+__all__ = ["apply_lora"]

comfy_diffusion/models.py

@@ -0,0 +1,245 @@
+"""Model management public API.
+
+This module must stay import-safe in CPU-only environments. It intentionally avoids
+importing ComfyUI loaders at module import time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from ._runtime import ensure_comfyui_on_path
+
+
+@dataclass
+class CheckpointResult:
+    """Container for objects produced by a ComfyUI checkpoint load."""
+
+    model: Any
+    clip: Any | None
+    vae: Any | None
+
+
+class ModelManager:
+    """Entry point for model-loading operations.
+
+    The implementation is intentionally deferred to later user stories; this class
+    exists now to provide a stable, side-effect-free import surface.
+    """
+
+    def __init__(self, models_dir: str | Path) -> None:
+        """Store and validate the models directory used by future load operations."""
+        path = Path(models_dir)
+
+        if not path.exists():
+            raise ValueError(f"models_dir does not exist: {path}")
+        if not path.is_dir():
+            raise ValueError(f"models_dir is not a directory: {path}")
+
+        self.models_dir = path
+
+        ensure_comfyui_on_path()
+        import folder_paths
+
+        folder_paths.add_model_folder_path(
+            "checkpoints", str(self.models_dir / "checkpoints"), is_default=True
+        )
+        folder_paths.add_model_folder_path(
+            "embeddings", str(self.models_dir / "embeddings"), is_default=True
+        )
+        folder_paths.add_model_folder_path(
+            "diffusion_models", str(self.models_dir / "unet"), is_default=True
+        )
+        folder_paths.add_model_folder_path(
+            "diffusion_models", str(self.models_dir / "diffusion_models"), is_default=False
+        )
+        folder_paths.add_model_folder_path(
+            "text_encoders", str(self.models_dir / "text_encoders"), is_default=True
+        )
+        folder_paths.add_model_folder_path(
+            "text_encoders", str(self.models_dir / "clip"), is_default=False
+        )
+        folder_paths.add_model_folder_path(
+            "vae", str(self.models_dir / "vae"), is_default=True
+        )
+
+    def load_checkpoint(self, filename: str) -> CheckpointResult:
+        """Load a checkpoint by filename from the configured checkpoints directory."""
+        ensure_comfyui_on_path()
+
+        requested_path = (self.models_dir / "checkpoints" / filename).resolve()
+        if not requested_path.is_file():
+            raise FileNotFoundError(f"checkpoint file not found: {requested_path}")
+
+        import folder_paths
+        from comfy import sd as comfy_sd
+
+        checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", filename)
+        loaded = comfy_sd.load_checkpoint_guess_config(
+            checkpoint_path,
+            output_vae=True,
+            output_clip=True,
+            embedding_directory=folder_paths.get_folder_paths("embeddings"),
+        )
+        model, clip, vae = loaded[:3]
+        return CheckpointResult(model=model, clip=clip, vae=vae)
+
+    def load_vae(self, path: str | Path) -> Any:
+        """Load a standalone VAE from a path or filename.
+
+        If ``path`` is an absolute path to an existing file, that file is loaded.
+        Otherwise ``path`` is treated as a filename under the ``vae`` folder.
+        """
+        ensure_comfyui_on_path()
+
+        import folder_paths
+        from comfy import sd as comfy_sd
+        from comfy import utils as comfy_utils
+
+        p = Path(path)
+        if p.is_absolute() and p.is_file():
+            vae_path = str(p.resolve())
+        elif p.is_absolute():
+            raise FileNotFoundError(f"vae file not found: {p}")
+        else:
+            name = path if isinstance(path, str) else p.name
+            vae_path = folder_paths.get_full_path_or_raise("vae", name)
+
+        state_dict, metadata = comfy_utils.load_torch_file(
+            vae_path, return_metadata=True
+        )
+        vae = comfy_sd.VAE(sd=state_dict, metadata=metadata)
+        vae.throw_exception_if_invalid()
+        return vae
+
+    def load_clip(
+        self,
+        path: str | Path,
+        *,
+        clip_type: str = "stable_diffusion",
+    ) -> Any:
+        """Load a standalone text encoder (CLIP) from a path or filename.
+
+        If ``path`` is an absolute path to an existing file, that file is loaded.
+        Otherwise ``path`` is treated as a filename under ``text_encoders`` / ``clip``.
+
+        ``clip_type`` selects the encoder architecture (e.g. ``"wan"`` for Wan / UMT5-XXL,
+        ``"stable_diffusion"``, ``"sd3"``, ``"flux"``). Must match the model weights.
+        """
+        ensure_comfyui_on_path()
+
+        import folder_paths
+        from comfy import sd as comfy_sd
+
+        p = Path(path)
+        if p.is_absolute() and p.is_file():
+            full_path = str(p.resolve())
+        elif p.is_absolute():
+            raise FileNotFoundError(f"clip file not found: {p}")
+        else:
+            name = path if isinstance(path, str) else p.name
+            full_path = folder_paths.get_full_path_or_raise("text_encoders", name)
+
+        clip_type_enum = getattr(
+            comfy_sd.CLIPType,
+            clip_type.upper(),
+            comfy_sd.CLIPType.STABLE_DIFFUSION,
+        )
+        return comfy_sd.load_clip(
+            ckpt_paths=[full_path],
+            embedding_directory=folder_paths.get_folder_paths("embeddings"),
+            clip_type=clip_type_enum,
+        )
+
+    def load_unet(self, path: str | Path) -> Any:
+        """Load a standalone diffusion model (UNet) from a path or filename.
+
+        If ``path`` is an absolute path to an existing file, that file is loaded.
+        Otherwise ``path`` is treated as a filename and resolved under the
+        ``diffusion_models`` / ``unet`` folders (see ComfyUI folder layout).
+        """
+        ensure_comfyui_on_path()
+
+        import folder_paths
+        from comfy import sd as comfy_sd
+
+        p = Path(path)
+        if p.is_absolute() and p.is_file():
+            full_path = str(p.resolve())
+        elif p.is_absolute():
+            raise FileNotFoundError(f"unet file not found: {p}")
+        else:
+            name = path if isinstance(path, str) else p.name
+            full_path = folder_paths.get_full_path_or_raise("diffusion_models", name)
+
+        return comfy_sd.load_diffusion_model(full_path)
+
+    def load_ltxv_audio_vae(self, path: str | Path) -> object:
+        """Load an LTXV audio VAE checkpoint from a path or filename.
+
+        If ``path`` is an absolute path to an existing file, that file is loaded.
+        Otherwise ``path`` is treated as a filename under the ``checkpoints`` folder.
+        """
+        ensure_comfyui_on_path()
+
+        import folder_paths
+        from comfy import utils as comfy_utils
+        from comfy.ldm.lightricks.vae.audio_vae import AudioVAE
+
+        p = Path(path)
+        if p.is_absolute() and p.is_file():
+            checkpoint_path = str(p.resolve())
+        elif p.is_absolute():
+            raise FileNotFoundError(f"ltxv audio vae file not found: {p}")
+        else:
+            name = path if isinstance(path, str) else p.name
+            checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", name)
+
+        state_dict, metadata = comfy_utils.load_torch_file(
+            checkpoint_path, return_metadata=True
+        )
+        return AudioVAE(state_dict, metadata)
+
+    def load_ltxav_text_encoder(
+        self, text_encoder_path: str | Path, checkpoint_path: str | Path
+    ) -> object:
+        """Load an LTXAV text encoder from two separate files.
+
+        ``text_encoder_path`` is the text encoder file (from ``text_encoders/``).
+        ``checkpoint_path`` is the companion checkpoint file (from ``checkpoints/``).
+        Both can be absolute paths to existing files or relative filenames resolved
+        via folder_paths.
+        """
+        ensure_comfyui_on_path()
+
+        import folder_paths
+        from comfy import sd as comfy_sd
+
+        te_p = Path(text_encoder_path)
+        if te_p.is_absolute() and te_p.is_file():
+            resolved_te = str(te_p.resolve())
+        elif te_p.is_absolute():
+            raise FileNotFoundError(f"ltxav text encoder file not found: {te_p}")
+        else:
+            name = text_encoder_path if isinstance(text_encoder_path, str) else te_p.name
+            resolved_te = folder_paths.get_full_path_or_raise("text_encoders", name)
+
+        ckpt_p = Path(checkpoint_path)
+        if ckpt_p.is_absolute() and ckpt_p.is_file():
+            resolved_ckpt = str(ckpt_p.resolve())
+        elif ckpt_p.is_absolute():
+            raise FileNotFoundError(f"ltxav checkpoint file not found: {ckpt_p}")
+        else:
+            name = checkpoint_path if isinstance(checkpoint_path, str) else ckpt_p.name
+            resolved_ckpt = folder_paths.get_full_path_or_raise("checkpoints", name)
+
+        return comfy_sd.load_clip(
+            ckpt_paths=[resolved_te, resolved_ckpt],
+            embedding_directory=folder_paths.get_folder_paths("embeddings"),
+            clip_type=comfy_sd.CLIPType.LTXV,
+        )
+
+
+__all__ = ["CheckpointResult", "ModelManager"]

comfy_diffusion/runtime.py

@@ -0,0 +1,86 @@
+"""Runtime diagnostics for comfy_diffusion."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from typing import Any
+
+
+def _python_version() -> str:
+    return ".".join(str(part) for part in sys.version_info[:3])
+
+
+def _runtime_not_found(python_version: str, detail: str = "") -> dict[str, Any]:
+    msg = "ComfyUI runtime not found. Run: git submodule update --init"
+    if detail:
+        msg += f" (or install missing deps). Cause: {detail}"
+    return {
+        "error": msg,
+        "comfyui_version": None,
+        "device": None,
+        "vram_total_mb": None,
+        "vram_free_mb": None,
+        "python_version": python_version,
+    }
+
+
+def _runtime_not_responsive(python_version: str, message: str) -> dict[str, Any]:
+    return {
+        "error": f"ComfyUI runtime is not responsive: {message}",
+        "comfyui_version": None,
+        "device": None,
+        "vram_total_mb": None,
+        "vram_free_mb": None,
+        "python_version": python_version,
+    }
+
+
+def _bytes_to_mb(value: int) -> int:
+    return value // (1024 * 1024)
+
+
+def check_runtime() -> dict[str, Any]:
+    """Return structured runtime diagnostics for the current Python process."""
+    python_version = _python_version()
+
+    from ._runtime import ensure_comfyui_on_path
+    ensure_comfyui_on_path()
+
+    try:
+        comfyui_version_module = importlib.import_module("comfyui_version")
+        model_management = importlib.import_module("comfy.model_management")
+    except Exception as exc:
+        return _runtime_not_found(python_version, str(exc))
+
+    try:
+        device = model_management.get_torch_device()
+    except Exception as exc:
+        return _runtime_not_responsive(python_version, str(exc))
+
+    comfyui_version = str(getattr(comfyui_version_module, "__version__", "unknown"))
+    device_str = str(device)
+    device_type = getattr(device, "type", "")
+
+    if device_type == "cpu" or device_str == "cpu":
+        return {
+            "comfyui_version": comfyui_version,
+            "device": "cpu",
+            "vram_total_mb": 0,
+            "vram_free_mb": 0,
+            "python_version": python_version,
+        }
+
+    try:
+        total_memory_bytes = model_management.get_total_memory(device)
+        free_memory_bytes = model_management.get_free_memory(device)
+    except Exception as exc:
+        return _runtime_not_responsive(python_version, str(exc))
+
+    return {
+        "comfyui_version": comfyui_version,
+        "device": device_str,
+        "vram_total_mb": _bytes_to_mb(int(total_memory_bytes)),
+        "vram_free_mb": _bytes_to_mb(int(free_memory_bytes)),
+        "python_version": python_version,
+    }