forced lora

handler.py

@@ -6,6 +6,7 @@ from io import BytesIO
 from pprint import pprint
 from typing import Any, Dict, List
 import os
+import re
 from pathlib import Path
 from typing import Union
 from concurrent.futures import ThreadPoolExecutor
@@ -87,6 +88,7 @@ class EndpointHandler:
         self.inference_progress = {}  # Dictionary to store progress of each request
         self.inference_images = {}  # Dictionary to store latest image of each request
         self.total_steps = {}
+        self.active_request_ids = set()
         self.inference_in_progress = False
 
         self.executor = ThreadPoolExecutor(
@@ -131,6 +133,18 @@ class EndpointHandler:
         self.pipe.enable_attention_slicing()
         # may need a requirement in the root with xformer
 
+        # Load loras one time only
+        # Must be replaced once we will know how to hot load/unload
+        # it use the own made load_lora function
+        self.load_selected_loras(
+            [
+                ["polyhedron_new_skin_v1.1", 0.2],
+                ["detailed_eye-10", 0.2],
+                ["add_detail", 0.3],
+                ["MuscleGirl_v1", 0.2],
+            ]
+        )
+
     def load_lora(self, pipeline, lora_path, lora_weight=0.5):
         state_dict = load_file(lora_path)
         LORA_PREFIX_UNET = "lora_unet"
@@ -218,10 +232,33 @@ class EndpointHandler:
         """Load Loras models, can lead to marvelous creations"""
         for model_name, weight in selections:
             lora_path = EndpointHandler.LORA_PATHS[model_name]
-            self.pipe = self.load_lora(
-                pipeline=self.pipe, lora_path=lora_path, lora_weight=weight
-            )
-        return self.pipe
+            # self.pipe.load_lora_weights(lora_path)
+            self.load_lora(self.pipe, lora_path, weight)
+
+    def clean_negative_prompt(self, negative_prompt):
+        """Clean negative prompt to remove already used negative prompt handlers"""
+
+        # negative_prompt = (
+        #     negative_prompt
+        #     + """, easynegative, badhandv4, bad-artist-anime, negfeetv2, ng_deepnegative_v1_75t, bad-hands-5, """
+        # )
+
+        tokens = [item["token"] for item in self.TEXTUAL_INVERSION]
+
+        # Retirer tous les tokens de negative_prompt s'ils existent déjà
+        for token in tokens:
+            # Utiliser une expression régulière pour un remplacement insensible à la casse
+            negative_prompt = re.sub(
+                r"\b" + re.escape(token) + r"\b",
+                "",
+                negative_prompt,
+                flags=re.IGNORECASE,
+            ).strip()
+
+        # Ajouter tous les tokens à la fin de negative_prompt
+        negative_prompt += " " + " ".join(tokens)
+
+        return negative_prompt
 
     def clean_request_data(self, request_id: str):
         """Clean up the data related to a specific request ID."""
@@ -235,6 +272,9 @@ class EndpointHandler:
         # Remove the request ID from the total_steps dictionary
         self.total_steps.pop(request_id, None)
 
+        # Delete request id
+        self.active_request_ids.discard(request_id)
+
         # Set inference to False
         self.inference_in_progress = False
 
@@ -349,17 +389,18 @@ class EndpointHandler:
         self.total_steps[request_id] = num_inference_steps
 
         # USe this to add automatically some negative prompts
-        forced_negative = (
-            negative_prompt
-            + """, easynegative, badhandv4, bad-artist-anime, negfeetv2, ng_deepnegative_v1_75t, bad-hands-5, """
-        )
+        forced_negative = self.clean_negative_prompt(negative_prompt)
 
         # Set the generator seed if provided
         generator = torch.Generator(device="cuda").manual_seed(seed) if seed else None
 
         # Load the provided Lora models
+        # self.pipe.unload_lora_weights()  # Unload models to avoid lora staking
         # if loras_model:
-        #     self.pipe = self.load_selected_loras(loras_model)
+        #     self.load_selected_loras(loras_model)
+
+        # set scale of loras, for now take only first scale of the loaded lora and apply to all until we find the way to apply specified scale
+        # scale = {"scale": loras_model[0][1]} if loras_model else None
 
         try:
             # 2. Process
@@ -376,8 +417,8 @@ class EndpointHandler:
                     callback=lambda step, timestep, latents: self.progress_callback(
                         step, timestep, latents, request_id, "progress"
                     ),
-                    callback_steps=5,  #  The frequency at which the callback function is called.
-                    # output_type="pt",
+                    callback_steps=5,
+                    # cross_attention_kwargs={"scale": 0.2},
                 ).images[0]
 
             # print(image)
@@ -405,6 +446,11 @@ class EndpointHandler:
             return {"flag": "error", "message": "Missing request_id."}
 
         if action == "check_progress":
+            if request_id not in self.active_request_ids:
+                return {
+                    "flag": "error",
+                    "message": "Request id doesn't match any active request.",
+                }
             return self.check_progress(request_id)
 
         elif action == "inference":
@@ -420,6 +466,7 @@ class EndpointHandler:
             self.inference_in_progress = True
             self.inference_progress[request_id] = 0
             self.inference_images[request_id] = None
+            self.active_request_ids.add(request_id)
 
             self.executor.submit(self.start_inference, data)