feat: device load

2025-11-02 22:27:24 +09:00
parent c96edd2c53
commit c87630816b
1 changed files with 3 additions and 1 deletions
@@ -2,6 +2,7 @@ import os

 import folder_paths  # pyright: ignore[reportMissingImports]
 import torch
+from comfy.model_management import get_torch_device  # pyright: ignore[reportMissingImports]
 from diffusers import StableDiffusionPipeline

 from ...domain.repositories import PipelineRepository
@@ -10,6 +11,7 @@ from ...domain.repositories import PipelineRepository
 class DiffusersPipelineRepository(PipelineRepository):
    def __init__(self) -> None:
        self.cache_dir = folder_paths.get_temp_directory()
+        self.device = get_torch_device()

    def convert_and_save_from_single_file(self, checkpoint_path: str, dtype: torch.dtype) -> str:
        checkpoint_name = os.path.basename(checkpoint_path)
@@ -29,5 +31,5 @@ class DiffusersPipelineRepository(PipelineRepository):
            pretrained_model_name_or_path=model_path,
            torch_dtype=dtype,
            cache_dir=self.cache_dir,
-        )
+        ).to(self.device)
        return pipe