refactor: architecture

This commit is contained in:
Limitex
2025-11-03 06:02:16 +09:00
parent 955fa8a7fa
commit 98b85781f3
3 changed files with 18 additions and 9 deletions
@@ -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 ---
+11 -3
View File
@@ -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
+3 -6
View File
@@ -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