Update handler.py

handler.py

@@ -207,6 +207,11 @@ class GenerationConfig:
     # Enhance-A-Video settings
     enable_enhance_a_video: bool = True
     enhance_a_video_weight: float = 4.0
+
+    # LoRA settings
+    lora_model_name: str = ""  # HuggingFace repo ID or path to LoRA model
+    lora_model_weight_file: str = ""  # Specific weight file to load from the LoRA model
+    lora_model_trigger: str = ""  # Optional trigger word to prepend to the prompt
     
     def validate_and_adjust(self) -> 'GenerationConfig':
         """Validate and adjust parameters to meet constraints"""
@@ -260,6 +265,9 @@ class EndpointHandler:
             torch_dtype=torch.bfloat16
         ).to("cuda")
 
+        # Initialize LoRA tracking
+        self._current_lora_model = None
+
         # Enable CPU offload for memory efficiency
         #self.text_to_video.enable_model_cpu_offload()
         #self.image_to_video.enable_model_cpu_offload()
@@ -393,7 +401,9 @@ class EndpointHandler:
                     - teacache_threshold (optional, float, default to 0.05): Amount of cache, 0 (original), 0.03 (1.6x speedup), 0.05 (Default, 2.1x speedup).
                     - enable_enhance_a_video (optional, bool, default to True): enable the enhance_a_video optimization
                     - enhance_a_video_weight(optional, float, default to 4.0): amount of video enhancement to apply
-            
+                    - lora_model_name(optional, str, default to ""): HuggingFace repo ID or path to LoRA model
+                    - lora_model_weight_file(optional, str, default to ""): Specific weight file to load from the LoRA model
+                    - lora_model_trigger(optional, str, default to ""): Optional trigger word to prepend to the prompt
         Returns:
             Dictionary containing:
                 - video: Base64 encoded MP4 data URI
@@ -450,6 +460,11 @@ class EndpointHandler:
             # Add enhance-a-video settings
             enable_enhance_a_video=params.get("enable_enhance_a_video", True),
             enhance_a_video_weight=params.get("enhance_a_video_weight", 4.0),
+
+            # LoRA settings
+            lora_model_name=params.get("lora_model_name", ""),
+            lora_model_weight_file=params.get("lora_model_weight_file", ""),
+            lora_model_trigger=params.get("lora_model_trigger", ""),
         ).validate_and_adjust()
         
         #logger.debug(f"Global request settings:")
@@ -494,6 +509,34 @@ class EndpointHandler:
                 }
                 #logger.info(f"Video model generation settings:")
                 #pprint.pprint(generation_kwargs)
+
+                # Handle LoRA loading/unloading
+                if hasattr(self, '_current_lora_model'):
+                    if self._current_lora_model != (config.lora_model_name, config.lora_model_weight_file):
+                        # Unload previous LoRA if it exists and is different
+                        if hasattr(self.text_to_video, 'unload_lora_weights'):
+                            self.text_to_video.unload_lora_weights()
+                        if hasattr(self.image_to_video, 'unload_lora_weights'):
+                            self.image_to_video.unload_lora_weights()
+    
+                if config.lora_model_name:
+                    # Load new LoRA
+                    if hasattr(self.text_to_video, 'load_lora_weights'):
+                        self.text_to_video.load_lora_weights(
+                            config.lora_model_name,
+                            weight_name=config.lora_model_weight_file if config.lora_model_weight_file else None
+                        )
+                    if hasattr(self.image_to_video, 'load_lora_weights'):
+                        self.image_to_video.load_lora_weights(
+                            config.lora_model_name,
+                            weight_name=config.lora_model_weight_file if config.lora_model_weight_file else None
+                        )
+                    self._current_lora_model = (config.lora_model_name, config.lora_model_weight_file)
+    
+                # Modify prompt if trigger word is provided
+                if config.lora_model_trigger:
+                    generation_kwargs["prompt"] = f"{config.lora_model_trigger} {generation_kwargs['prompt']}"
+    
                 
                 # Check if image-to-video generation is requested
                 if input_image: