diff --git a/src/domain/model/_lcm_lora.py b/src/domain/model/_lcm_lora.py index af42755..ad97092 100644 --- a/src/domain/model/_lcm_lora.py +++ b/src/domain/model/_lcm_lora.py @@ -1,12 +1,13 @@ from dataclasses import dataclass -from typing import Any + +import torch @dataclass(frozen=True) class LcmLora: """LCM LoRA weights for Stream Diffusion.""" - weights: dict[str, Any] + weights: dict[str, torch.Tensor] def __post_init__(self) -> None: if not isinstance(self.weights, dict): diff --git a/src/domain/repositories/_lcm_lora_repository.py b/src/domain/repositories/_lcm_lora_repository.py index 51739ef..5bd35ba 100644 --- a/src/domain/repositories/_lcm_lora_repository.py +++ b/src/domain/repositories/_lcm_lora_repository.py @@ -1,10 +1,11 @@ from abc import ABC, abstractmethod -from typing import Any + +import torch class LcmLoraRepository(ABC): @abstractmethod - def load_lcm_lora(self, lora_path: str) -> dict[str, Any]: + def load_lcm_lora(self, lora_path: str) -> dict[str, torch.Tensor]: """Load LCM LoRA weights from file. Args: diff --git a/src/infra/diffusers/_lcm_lora_repository.py b/src/infra/diffusers/_lcm_lora_repository.py index 44654a9..36738fd 100644 --- a/src/infra/diffusers/_lcm_lora_repository.py +++ b/src/infra/diffusers/_lcm_lora_repository.py @@ -1,12 +1,11 @@ -from typing import Any - +import torch from safetensors.torch import load_file from ...domain.repositories import LcmLoraRepository class DiffusersLcmLoraRepository(LcmLoraRepository): - def load_lcm_lora(self, lora_path: str) -> dict[str, Any]: + def load_lcm_lora(self, lora_path: str) -> dict[str, torch.Tensor]: """Load LCM LoRA weights from safetensors file. Args: diff --git a/src/nodes/dto/_lcm_lora_dto.py b/src/nodes/dto/_lcm_lora_dto.py index 0b1b9c0..26e195b 100644 --- a/src/nodes/dto/_lcm_lora_dto.py +++ b/src/nodes/dto/_lcm_lora_dto.py @@ -1,5 +1,7 @@ from dataclasses import dataclass -from typing import Any, ClassVar +from typing import ClassVar + +import torch from ...domain.model import LcmLora @@ -17,7 +19,7 @@ class ComfyUILcmLoraDTO: COMFY_TYPE: ClassVar[str] = "DIFFUSERS_LCM_LORA" - weights: dict[str, Any] + weights: dict[str, torch.Tensor] @classmethod def from_domain(cls, lcm_lora: LcmLora) -> "ComfyUILcmLoraDTO":