From c87630816bb5bcaf2f814bf158aba7d200710c26 Mon Sep 17 00:00:00 2001
From: Limitex <76650151+Limitex@users.noreply.github.com>
Date: Sun, 2 Nov 2025 22:00:59 +0900
Subject: [PATCH] feat: device load

---
 src/infra/diffusers/_pipeline_repository.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/infra/diffusers/_pipeline_repository.py b/src/infra/diffusers/_pipeline_repository.py
index dde10c0..a024f75 100644
--- a/src/infra/diffusers/_pipeline_repository.py
+++ b/src/infra/diffusers/_pipeline_repository.py
@@ -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