refactor: diffusers impl

2025-11-23 10:02:37 +09:00
parent 3dbb83d209
commit e17bdfc621
2 changed files with 9 additions and 3 deletions
@@ -36,4 +36,10 @@ class DiffusersPipelineRepository(PipelineRepository):
            torch_dtype=dtype,
            cache_dir=self.cache_dir,
        ).to(self.device)
+        pipe.safety_checker = (  # type: ignore[attr-defined]
+            None
+            if pipe.safety_checker is None  # type: ignore[attr-defined]
+            else lambda images, **_kwargs: (images, [False])
+        )
+        pipe.enable_attention_slicing()  # type: ignore[attr-defined]
        return pipe
@@ -24,10 +24,10 @@ class DiffusersSamplerRepository(SamplerRepository):
        cfg: CFGScale,
        seed: Seed,
    ) -> list[Image.Image]:
-        result = pipeline(  # type: ignore[operator]
+        pipeline.vae = vae  # type: ignore[attr-defined]
+        pipeline.scheduler = scheduler  # type: ignore[attr-defined]
+        result = pipeline.to(self.device)(  # type: ignore[attr-defined]
            prompt_embeds=positive_embeds,
-            vae=vae,
-            scheduler=scheduler,
            height=image_size.height,
            width=image_size.width,
            num_inference_steps=steps.value,