Create LiteVisionPipeline.py

LiteVisionPipeline.py

@@ -0,0 +1,44 @@
+import torch
+from diffusers import StableDiffusionPipeline, LCMScheduler
+
+
+class LiteVisionPipeline(StableDiffusionPipeline):
+    """
+    LiteVisionPipeline v1
+    - Custom pipeline for LiteVision models
+    - LCM scheduler preset
+    - Low VRAM & fast inference defaults
+    - Compatible with DiffusionPipeline ecosystem
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # ---- Scheduler preset (LCM) ----
+        self.scheduler = LCMScheduler.from_config(self.scheduler.config)
+
+        # ---- Memory optimizations ----
+        self.enable_attention_slicing()
+        self.enable_vae_slicing()
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt,
+        negative_prompt=None,
+        num_inference_steps: int = 6,
+        guidance_scale: float = 1.5,
+        **kwargs,
+    ):
+        """
+        LiteVision default generation call
+        - Optimized for LCM-style inference
+        """
+
+        return super().__call__(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            **kwargs,
+        )