Update services/generation.py

services/generation.py

@@ -1,13 +1,14 @@
 # services/generation.py
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from diffusers import StableDiffusionPipeline, DiffusionPipeline, DPMSolverMultistepScheduler
 from PIL import Image
 import config
 from utils.helpers import decode_base64_image, encode_image_base64, encode_video_base64
 import logging
 import gc # Garbage collector
 from typing import List 
+from diffusers import StableDiffusionPipeline, DiffusionPipeline, DPMSolverMultistepScheduler, LCMScheduler # Import LCMScheduler
+from peft import PeftConfig # Import PeftConfig (if needed, usually handled by load_lora_weights)
 
 logger = logging.getLogger(__name__)
 
@@ -27,27 +28,43 @@ def load_models():
         model_cache["text_model"] = AutoModelForSeq2SeqLM.from_pretrained(config.TEXT_MODEL_NAME).to(config.DEVICE)
         logger.info("Text model loaded.")
 
-        # Image Generation Model
-        logger.info(f"Loading image model: {config.IMAGE_MODEL_NAME}")
-        image_pipeline = StableDiffusionPipeline.from_pretrained(
-            config.IMAGE_MODEL_NAME,
-            torch_dtype=config.DTYPE
-        )
-        # Optimization: Use a faster scheduler
+    # --- Image Generation Model ---
+    logger.info(f"Loading image model: {config.IMAGE_MODEL_NAME}")
+    image_pipeline = StableDiffusionPipeline.from_pretrained(
+        config.IMAGE_MODEL_NAME,
+        torch_dtype=config.DTYPE
+    )
+
+    # --- Load LCM LoRA ---
+    try:
+        logger.info(f"Loading LCM LoRA: {config.IMAGE_LCM_LORA_NAME}")
+        # Load LoRA weights directly into the pipeline
+        image_pipeline.load_lora_weights(config.IMAGE_LCM_LORA_NAME)
+        # Fuse LoRA for potential speedup (optional, test impact)
+        # image_pipeline.fuse_lora()
+        logger.info("LCM LoRA loaded successfully.")
+
+        # --- IMPORTANT: Set LCM Scheduler ---
+        image_pipeline.scheduler = LCMScheduler.from_config(image_pipeline.scheduler.config)
+        logger.info("Switched scheduler to LCMScheduler.")
+
+    except Exception as e:
+        logger.warning(f"Could not load or apply LCM LoRA '{config.IMAGE_LCM_LORA_NAME}'. Falling back to base model scheduler. Error: {e}", exc_info=True)
+        # Fallback to a standard fast scheduler if LCM fails
         image_pipeline.scheduler = DPMSolverMultistepScheduler.from_config(image_pipeline.scheduler.config)
-        image_pipeline = image_pipeline.to(config.DEVICE)
-        # Optional: Enable attention slicing for lower VRAM usage on GPU
-        if config.DEVICE == "cuda":
-             try:
-                 # Requires pip install xformers - uncomment if installed
-                 # image_pipeline.enable_xformers_memory_efficient_attention()
-                 pass # Use default if xformers not installed/wanted
-             except ImportError:
-                 logger.warning("xformers not installed. Memory efficient attention not enabled.")
-             # image_pipeline.enable_attention_slicing() # Alternative if xformers not available
-
-        model_cache["image_pipeline"] = image_pipeline
-        logger.info("Image model loaded.")
+
+
+    image_pipeline = image_pipeline.to(config.DEVICE)
+    if config.DEVICE == "cuda":
+         try:
+             # image_pipeline.enable_xformers_memory_efficient_attention()
+             pass
+         except ImportError:
+             logger.warning("xformers not installed...")
+         # image_pipeline.enable_attention_slicing()
+
+    model_cache["image_pipeline"] = image_pipeline
+    logger.info("Image model setup complete.")
 
 
         # Video Generation Model