From 98b85781f3ba2b1632f0ffd019f48767f3c87c86 Mon Sep 17 00:00:00 2001
From: Limitex <76650151+Limitex@users.noreply.github.com>
Date: Mon, 3 Nov 2025 05:47:42 +0900
Subject: [PATCH] refactor: architecture

---
 src/infra/diffusers/_autoencoder_repository.py |  4 ++++
 src/service/_pipeline_service.py               | 14 +++++++++++---
 src/ui/_pipeline_handler.py                    |  9 +++------
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/infra/diffusers/_autoencoder_repository.py b/src/infra/diffusers/_autoencoder_repository.py
index b37219f..c113703 100644
--- a/src/infra/diffusers/_autoencoder_repository.py
+++ b/src/infra/diffusers/_autoencoder_repository.py
@@ -32,11 +32,15 @@ class DiffusersAutoencoderRepository(AutoencoderRepository):
         return ckpt_cache_path
 
     def load_autoencoder_from_path(self, model_path: str, dtype: torch.dtype) -> AutoencoderKL:
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Model path not found: {model_path}")
         vae: AutoencoderKL = AutoencoderKL.from_pretrained(  # type: ignore[no-untyped-call]
             pretrained_model_name_or_path=model_path,
             torch_dtype=dtype,
             cache_dir=self.cache_dir,
         )
+        if vae is None:
+            raise RuntimeError(f"Failed to load AutoencoderKL from path: {model_path}")
         return vae
 
     # --- TECHNICAL DEBT & RISK WARNING ---
diff --git a/src/service/_pipeline_service.py b/src/service/_pipeline_service.py
index 83af727..1ca6c27 100644
--- a/src/service/_pipeline_service.py
+++ b/src/service/_pipeline_service.py
@@ -2,7 +2,7 @@ import os
 
 import torch
 
-from ..domain.model import Pipeline
+from ..domain.model import Clip, Pipeline
 from ..domain.repositories import PipelineRepository
 
 
@@ -11,7 +11,7 @@ class PipelineService:
         self.pipeline_repo = pipeline_repo
         self.dtype = torch.float32
 
-    def create_from_checkpoint(self, checkpoint_full_path: str) -> Pipeline:
+    def create_from_checkpoint(self, checkpoint_full_path: str) -> tuple[Pipeline, Clip]:
         path = self.pipeline_repo.convert_and_save_from_single_file(
             checkpoint_full_path, self.dtype
         )
@@ -22,4 +22,12 @@ class PipelineService:
         if pipe is None:
             raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_full_path}")
 
-        return Pipeline(pipeline=pipe, path=path)
+        pipeline = Pipeline(pipeline=pipe, path=path)
+
+        clip = Clip(
+            tokenizer=pipe.tokenizer,  # type: ignore[attr-defined]
+            text_encoder=pipe.text_encoder,  # type: ignore[attr-defined]
+            path=checkpoint_full_path,
+        )
+
+        return pipeline, clip
diff --git a/src/ui/_pipeline_handler.py b/src/ui/_pipeline_handler.py
index 0a451a6..1e123df 100644
--- a/src/ui/_pipeline_handler.py
+++ b/src/ui/_pipeline_handler.py
@@ -11,12 +11,9 @@ class PipelineHandler:
     def create(self, checkpoint_path: str) -> tuple[Pipeline, Clip]:
         if not os.path.exists(checkpoint_path):
             raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
-        pipeline = self.create_pipeline_service.create_from_checkpoint(checkpoint_path)
+        pipeline, clip = self.create_pipeline_service.create_from_checkpoint(checkpoint_path)
         if pipeline is None:
             raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_path}")
-        clip = Clip(
-            tokenizer=pipeline.pipeline.tokenizer,  # type: ignore[attr-defined]
-            text_encoder=pipeline.pipeline.text_encoder,  # type: ignore[attr-defined]
-            path=checkpoint_path,
-        )
+        if clip is None:
+            raise RuntimeError(f"Failed to extract clip from checkpoint: {checkpoint_path}")
         return pipeline, clip