Merge remote-tracking branch 'origin/main' into packing

Author: OleehyO

src/cogkit/api/services/image_generation.py

@@ -58,15 +58,17 @@ def generate(
         if model not in self._models:
             raise ValueError(f"Model {model} not loaded")
         width, height = list(map(int, size.split("x")))
-
         # TODO: Refactor this to switch by LoRA endpoint API
         if lora_path != self._current_lora[model]:
             if lora_path is not None:
                 adapter_name = os.path.basename(lora_path)
-                _logger.info(f"Loading LORA weights from {adapter_name}")
+                _logger.info(
+                    f"Loading LORA weights from {adapter_name} and unload previous weights {self._current_lora[model]}"
+                )
+                unload_lora_checkpoint(self._models[model])
                 load_lora_checkpoint(self._models[model], lora_path, lora_scale)
             else:
-                _logger.info("Unloading LORA weights")
+                _logger.info(f"Unloading LORA weights {self._current_lora[model]}")
                 unload_lora_checkpoint(self._models[model])
 
             self._current_lora[model] = lora_path

src/cogkit/utils/lora.py

@@ -9,8 +9,8 @@ def load_lora_checkpoint(
     lora_model_id_or_path: str,
     lora_scale: float = 1.0,
 ) -> None:
-    pipeline.load_lora_weights(lora_model_id_or_path)
-    pipeline.fuse_lora(components=["transformer"], lora_scale=lora_scale)
+    pipeline.load_lora_weights(lora_model_id_or_path, lora_scale=lora_scale)
+    # pipeline.fuse_lora(components=["transformer"], lora_scale=lora_scale)
 
 
 def unload_lora_checkpoint(

web/infer.py

@@ -59,11 +59,7 @@ def get_lora_paths():
     if not os.path.exists(lora_dir):
         os.makedirs(lora_dir, exist_ok=True)
         return ["None"]
-    checkpoint_dirs = [
-        d
-        for d in os.listdir(lora_dir)
-        if os.path.isdir(os.path.join(lora_dir, d)) and d.startswith("checkpoint")
-    ]
+    checkpoint_dirs = [d for d in os.listdir(lora_dir) if os.path.isdir(os.path.join(lora_dir, d))]
 
     if not checkpoint_dirs:
         return ["None"]
@@ -206,7 +202,7 @@ def refresh_lora_paths():
                         minimum=1,
                         maximum=50,
                         step=1,
-                        value=1,
+                        value=50,
                     )
                     guidance_scale = gr.Slider(
                         label="Guidance Scale",