Make lora loading api endpoint functional

api.py

@@ -1,17 +1,38 @@
-from typing import Optional
+from typing import Literal, Optional, TYPE_CHECKING
 
 import numpy as np
 from fastapi import FastAPI
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel, Field
 from platform import system
 
+if TYPE_CHECKING:
+    from flux_pipeline import FluxPipeline
+
 if system() == "Windows":
     MAX_RAND = 2**16 - 1
 else:
     MAX_RAND = 2**32 - 1
 
-app = FastAPI()
+
+class AppState:
+    model: "FluxPipeline"
+
+
+class FastAPIApp(FastAPI):
+    state: AppState
+
+
+class LoraArgs(BaseModel):
+    scale: Optional[float] = 1.0
+    path: Optional[str] = None
+    name: Optional[str] = None
+    action: Optional[Literal["load", "unload"]] = "load"
+
+
+class LoraLoadResponse(BaseModel):
+    status: Literal["success", "error"]
+    message: Optional[str] = None
 
 
 class GenerateArgs(BaseModel):
@@ -27,7 +48,75 @@ class GenerateArgs(BaseModel):
     init_image: Optional[str] = None
 
 
+app = FastAPIApp()
+
+
 @app.post("/generate")
 def generate(args: GenerateArgs):
+    """
+    Generates an image from the Flux flow transformer.
+
+    Args:
+        args (GenerateArgs): Arguments for image generation:
+
+            - `prompt`: The prompt used for image generation.
+
+            - `width`: The width of the image.
+
+            - `height`: The height of the image.
+
+            - `num_steps`: The number of steps for the image generation.
+
+            - `guidance`: The guidance for image generation, represents the
+                influence of the prompt on the image generation.
+
+            - `seed`: The seed for the image generation.
+
+            - `strength`: strength for image generation, 0.0 - 1.0.
+                Represents the percent of diffusion steps to run,
+                setting the init_image as the noised latent at the
+                given number of steps.
+
+            - `init_image`: Base64 encoded image or path to image to use as the init image.
+
+    Returns:
+        StreamingResponse: The generated image as streaming jpeg bytes.
+    """
     result = app.state.model.generate(**args.model_dump())
     return StreamingResponse(result, media_type="image/jpeg")
+
+
+@app.post("/lora", response_model=LoraLoadResponse)
+def lora_action(args: LoraArgs):
+    """
+    Loads or unloads a LoRA checkpoint into / from the Flux flow transformer.
+
+    Args:
+        args (LoraArgs): Arguments for the LoRA action:
+
+            - `scale`: The scaling factor for the LoRA weights.
+            - `path`: The path to the LoRA checkpoint.
+            - `name`: The name of the LoRA checkpoint.
+            - `action`: The action to perform, either "load" or "unload".
+
+    Returns:
+        LoraLoadResponse: The status of the LoRA action.
+    """
+    try:
+        if args.action == "load":
+            app.state.model.load_lora(args.path, args.scale, args.name)
+        elif args.action == "unload":
+            app.state.model.unload_lora(args.name if args.name else args.path)
+        else:
+            return JSONResponse(
+                content={
+                    "status": "error",
+                    "message": f"Invalid action, expected 'load' or 'unload', got {args.action}",
+                },
+                status_code=400,
+            )
+    except Exception as e:
+        return JSONResponse(
+            status_code=500, content={"status": "error", "message": str(e)}
+        )
+    return JSONResponse(status_code=200, content={"status": "success"})

flux_pipeline.py

@@ -2,7 +2,7 @@ import io
 import math
 import random
 import warnings
-from typing import TYPE_CHECKING, Callable, List, OrderedDict, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, OrderedDict, Union
 
 import numpy as np
 from PIL import Image
@@ -149,7 +149,10 @@ class FluxPipeline:
         return cuda_generator, seed
 
     def load_lora(
-        self, lora_path: Union[str, OrderedDict[str, torch.Tensor]], scale: float
+        self,
+        lora_path: Union[str, OrderedDict[str, torch.Tensor]],
+        scale: float,
+        name: Optional[str] = None,
     ):
         """
         Loads a LoRA checkpoint into the Flux flow transformer.
@@ -160,9 +163,9 @@ class FluxPipeline:
         Args:
             lora_path (str | OrderedDict[str, torch.Tensor]): Path to the LoRA checkpoint or an ordered dictionary containing the LoRA weights.
             scale (float): Scaling factor for the LoRA weights.
-
+            name (str): Name of the LoRA checkpoint, optionally can be left as None, since it only acts as an identifier.
         """
-        self.model.load_lora(lora_path, scale)
+        self.model.load_lora(path=lora_path, scale=scale, name=name)
 
     def unload_lora(self, path_or_identifier: str):
         """
@@ -171,7 +174,7 @@ class FluxPipeline:
         Args:
             path_or_identifier (str): Path to the LoRA checkpoint or the name given to the LoRA checkpoint when it was loaded.
         """
-        self.model.unload_lora(path_or_identifier)
+        self.model.unload_lora(path_or_identifier=path_or_identifier)
 
     @torch.inference_mode()
     def compile(self):