type hints and Pydantic

server/connection_manager.py

@@ -1,12 +1,12 @@
-from typing import Dict, Union
 from uuid import UUID
 import asyncio
 from fastapi import WebSocket
 from starlette.websockets import WebSocketState
 import logging
-from types import SimpleNamespace
+from typing import Any, TypeVar
+from util import ParamsModel
 
-Connections = Dict[UUID, Dict[str, Union[WebSocket, asyncio.Queue]]]
+Connections = dict[UUID, dict[str, WebSocket | asyncio.Queue]]
 
 
 class ServerFullException(Exception):
@@ -44,13 +44,13 @@ class ConnectionManager:
     def check_user(self, user_id: UUID) -> bool:
         return user_id in self.active_connections
 
-    async def update_data(self, user_id: UUID, new_data: SimpleNamespace):
+    async def update_data(self, user_id: UUID, new_data: ParamsModel):
         user_session = self.active_connections.get(user_id)
         if user_session:
             queue = user_session["queue"]
             await queue.put(new_data)
 
-    async def get_latest_data(self, user_id: UUID) -> SimpleNamespace:
+    async def get_latest_data(self, user_id: UUID) -> ParamsModel | None:
         user_session = self.active_connections.get(user_id)
         if user_session:
             queue = user_session["queue"]
@@ -58,6 +58,7 @@ class ConnectionManager:
                 return await queue.get()
             except asyncio.QueueEmpty:
                 return None
+        return None
 
     def delete_user(self, user_id: UUID):
         user_session = self.active_connections.pop(user_id, None)
@@ -86,7 +87,7 @@ class ConnectionManager:
             await websocket.close()
         self.delete_user(user_id)
 
-    async def send_json(self, user_id: UUID, data: Dict):
+    async def send_json(self, user_id: UUID, data: dict):
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
@@ -94,18 +95,22 @@ class ConnectionManager:
         except Exception as e:
             logging.error(f"Error: Send json: {e}")
 
-    async def receive_json(self, user_id: UUID) -> Dict:
+    async def receive_json(self, user_id: UUID) -> dict | None:
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
                 return await websocket.receive_json()
+            return None
         except Exception as e:
             logging.error(f"Error: Receive json: {e}")
+            return None
 
-    async def receive_bytes(self, user_id: UUID) -> bytes:
+    async def receive_bytes(self, user_id: UUID) -> bytes | None:
         try:
             websocket = self.get_websocket(user_id)
             if websocket:
                 return await websocket.receive_bytes()
+            return None
         except Exception as e:
             logging.error(f"Error: Receive bytes: {e}")
+            return None

server/main.py

@@ -10,30 +10,53 @@ import logging
 from config import config, Args
 from connection_manager import ConnectionManager, ServerFullException
 import uuid
+from uuid import UUID
 import time
-from types import SimpleNamespace
-from util import pil_to_frame, bytes_to_pil, is_firefox, get_pipeline_class
+from typing import Any, Protocol, TypeVar, runtime_checkable, cast
+from util import pil_to_frame, bytes_to_pil, is_firefox, get_pipeline_class, ParamsModel
 from device import device, torch_dtype
 import asyncio
 import os
 import time
 import torch
+from pydantic import BaseModel, create_model
+
+@runtime_checkable
+class BasePipeline(Protocol):
+    class Info:
+        @classmethod
+        def schema(cls) -> dict[str, Any]:
+            ...
+        page_content: str | None
+        input_mode: str
+    
+    class InputParams(ParamsModel):
+        @classmethod
+        def schema(cls) -> dict[str, Any]:
+            ...
+        
+        def dict(self) -> dict[str, Any]:
+            ...
+            
+    def predict(self, params: ParamsModel) -> Image.Image | None:
+        ...
 
 
 THROTTLE = 1.0 / 120
 
 
 class App:
-    def __init__(self, config: Args, pipeline):
+    def __init__(self, config: Args, pipeline_instance: BasePipeline):
         self.args = config
-        self.pipeline = pipeline
+        self.pipeline = pipeline_instance
         self.app = FastAPI()
         self.conn_manager = ConnectionManager()
+        self.safety_checker: SafetyChecker | None = None
         if self.args.safety_checker:
             self.safety_checker = SafetyChecker(device=device.type)
         self.init_app()
 
-    def init_app(self):
+    def init_app(self) -> None:
         self.app.add_middleware(
             CORSMiddleware,
             allow_origins=["*"],
@@ -43,7 +66,7 @@ class App:
         )
 
         @self.app.websocket("/api/ws/{user_id}")
-        async def websocket_endpoint(user_id: uuid.UUID, websocket: WebSocket):
+        async def websocket_endpoint(user_id: UUID, websocket: WebSocket) -> None:
             try:
                 await self.conn_manager.connect(
                     user_id, websocket, self.args.max_queue_size
@@ -55,9 +78,9 @@ class App:
                 await self.conn_manager.disconnect(user_id)
                 logging.info(f"User disconnected: {user_id}")
 
-        async def handle_websocket_data(user_id: uuid.UUID):
+        async def handle_websocket_data(user_id: UUID) -> None:
             if not self.conn_manager.check_user(user_id):
-                return HTTPException(status_code=404, detail="User not found")
+                raise HTTPException(status_code=404, detail="User not found")
             last_time = time.time()
             try:
                 while True:
@@ -75,19 +98,29 @@ class App:
                         await self.conn_manager.disconnect(user_id)
                         return
                     data = await self.conn_manager.receive_json(user_id)
+                    if data is None:
+                        continue
+                        
                     if data["status"] == "next_frame":
-                        info = pipeline.Info()
-                        params = await self.conn_manager.receive_json(user_id)
-                        params = pipeline.InputParams(**params)
-                        params = SimpleNamespace(**params.dict())
+                        info = self.pipeline.Info()
+                        params_data = await self.conn_manager.receive_json(user_id)
+                        if params_data is None:
+                            continue
+                            
+                        params = self.pipeline.InputParams.model_validate(params_data)
+                        
                         if info.input_mode == "image":
                             image_data = await self.conn_manager.receive_bytes(user_id)
-                            if len(image_data) == 0:
+                            if image_data is None or len(image_data) == 0:
                                 await self.conn_manager.send_json(
                                     user_id, {"status": "send_frame"}
                                 )
                                 continue
-                            params.image = bytes_to_pil(image_data)
+                            
+                            # Create a new Pydantic model with the image field
+                            params_dict = params.model_dump()
+                            params_dict["image"] = bytes_to_pil(image_data)
+                            params = self.pipeline.InputParams.model_validate(params_dict)
 
                         await self.conn_manager.update_data(user_id, params)
                         await self.conn_manager.send_json(user_id, {"status": "wait"})
@@ -97,29 +130,32 @@ class App:
                 await self.conn_manager.disconnect(user_id)
 
         @self.app.get("/api/queue")
-        async def get_queue_size():
+        async def get_queue_size() -> JSONResponse:
             queue_size = self.conn_manager.get_user_count()
             return JSONResponse({"queue_size": queue_size})
 
         @self.app.get("/api/stream/{user_id}")
-        async def stream(user_id: uuid.UUID, request: Request):
+        async def stream(user_id: UUID, request: Request) -> StreamingResponse:
             try:
-
-                async def generate():
-                    last_params = SimpleNamespace()
+                async def generate() -> bytes:
+                    last_params: ParamsModel | None = None
                     while True:
                         last_time = time.time()
                         await self.conn_manager.send_json(
                             user_id, {"status": "send_frame"}
                         )
                         params = await self.conn_manager.get_latest_data(user_id)
-                        if params.__dict__ == last_params.__dict__ or params is None:
+                        
+                        if (params is None or 
+                            (last_params is not None and 
+                             params.model_dump() == last_params.model_dump())):
                             await asyncio.sleep(THROTTLE)
                             continue
-                        last_params: SimpleNamespace = params
-                        image = pipeline.predict(params)
+                            
+                        last_params = params
+                        image = self.pipeline.predict(params)
 
-                        if self.args.safety_checker:
+                        if self.args.safety_checker and self.safety_checker is not None and image is not None:
                             image, has_nsfw_concept = self.safety_checker(image)
                             if has_nsfw_concept:
                                 image = None
@@ -141,23 +177,24 @@ class App:
                 )
             except Exception as e:
                 logging.error(f"Streaming Error: {e}, {user_id} ")
-                return HTTPException(status_code=404, detail="User not found")
+                raise HTTPException(status_code=404, detail="User not found")
 
         # route to setup frontend
         @self.app.get("/api/settings")
-        async def settings():
-            info_schema = pipeline.Info.schema()
-            info = pipeline.Info()
-            if info.page_content:
+        async def settings() -> JSONResponse:
+            info_schema = self.pipeline.Info.schema()
+            info = self.pipeline.Info()
+            page_content = ""
+            if hasattr(info, 'page_content') and info.page_content:
                 page_content = markdown2.markdown(info.page_content)
 
-            input_params = pipeline.InputParams.schema()
+            input_params = self.pipeline.InputParams.schema()
             return JSONResponse(
                 {
                     "info": info_schema,
                     "input_params": input_params,
                     "max_queue_size": self.args.max_queue_size,
-                    "page_content": page_content if info.page_content else "",
+                    "page_content": page_content,
                 }
             )
 
@@ -169,17 +206,35 @@ class App:
         )
 
 
+# def create_app(config):
+#     print(f"Device: {device}")
+#     print(f"torch_dtype: {torch_dtype}")
+
+#     # Create pipeline once
+#     pipeline_class = get_pipeline_class(config.pipeline)
+#     pipeline_instance = pipeline_class(config, device, torch_dtype)
+
+#     # Pass the existing pipeline instance to App
+#     app = App(config, pipeline_instance).app
+#     return app
+
+
+# Create app instance at module level
 print(f"Device: {device}")
 print(f"torch_dtype: {torch_dtype}")
+
 pipeline_class = get_pipeline_class(config.pipeline)
-pipeline = pipeline_class(config, device, torch_dtype)
-app = App(config, pipeline).app
+pipeline_instance = pipeline_class(config, device, torch_dtype)
+app = App(config, pipeline_instance).app  # This creates the FastAPI app instance
+
 
 if __name__ == "__main__":
     import uvicorn
 
+    # app = create_app(config)  # Create the app once
+
     uvicorn.run(
-        "main:app",
+        app,
         host=config.host,
         port=config.port,
         reload=config.reload,

server/pipelines/img2imgFlux.py

@@ -0,0 +1,157 @@
+import torch
+
+from optimum.quanto import freeze, qfloat8, quantize
+from transformers.modeling_utils import PreTrainedModel
+
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    AutoencoderKL,
+    AutoencoderTiny,
+    FluxImg2ImgPipeline,
+    FluxPipeline,
+)
+
+from diffusers import (
+    FluxImg2ImgPipeline,
+    FluxPipeline,
+    FluxTransformer2DModel,
+    GGUFQuantizationConfig,
+)
+
+try:
+    import intel_extension_for_pytorch as ipex  # type: ignore
+except:
+    pass
+
+import psutil
+from config import Args
+from pydantic import BaseModel, Field
+from PIL import Image
+from pathlib import Path
+import math
+import gc
+
+
+# model_path = "black-forest-labs/FLUX.1-dev"
+model_path = "black-forest-labs/FLUX.1-schnell"
+base_model_path = "black-forest-labs/FLUX.1-schnell"
+taesd_path = "madebyollin/taef1"
+subfolder = "transformer"
+transformer_path = model_path
+models_path = Path("models")
+
+default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
+default_negative_prompt = "blurry, low quality, render, 3D, oversaturated"
+page_content = """
+<h1 class="text-3xl font-bold">Real-Time FLUX</h1>
+
+"""
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class Pipeline:
+    class Info(BaseModel):
+        name: str = "img2img"
+        title: str = "Image-to-Image SDXL"
+        description: str = "Generates an image from a text prompt"
+        input_mode: str = "image"
+        page_content: str = page_content
+
+    class InputParams(BaseModel):
+        prompt: str = Field(
+            default_prompt,
+            title="Prompt",
+            field="textarea",
+            id="prompt",
+        )
+        seed: int = Field(
+            2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
+        )
+        steps: int = Field(
+            1, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
+        )
+        width: int = Field(
+            256, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+        )
+        height: int = Field(
+            256, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+        )
+        strength: float = Field(
+            0.5,
+            min=0.25,
+            max=1.0,
+            step=0.001,
+            title="Strength",
+            field="range",
+            hide=True,
+            id="strength",
+        )
+        guidance: float = Field(
+            3.5,
+            min=0,
+            max=20,
+            step=0.001,
+            title="Guidance",
+            hide=True,
+            field="range",
+            id="guidance",
+        )
+
+    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
+        # ckpt_path = (
+        #     "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
+        # )
+        print("Loading model")
+        # ckpt_path: str = "https://huggingface.co/city96/FLUX.1-schnell-gguf/blob/main/flux1-schnell-Q6_K.gguf"
+        ckpt_path: str = "https://huggingface.co/city96/FLUX.1-schnell-gguf/blob/main/flux1-schnell-Q4_K_S.gguf"
+        transformer = FluxTransformer2DModel.from_single_file(
+            ckpt_path,
+            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
+            torch_dtype=torch.bfloat16,
+        )
+
+        # else:
+        pipe = FluxImg2ImgPipeline.from_pretrained(
+            # "black-forest-labs/FLUX.1-dev",
+            "black-forest-labs/FLUX.1-Schnell",
+            transformer=transformer,
+            torch_dtype=torch.bfloat16,
+        )
+        if args.taesd:
+            pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_path, torch_dtype=torch.bfloat16, use_safetensors=True
+            )
+        # pipe.enable_model_cpu_offload()
+        pipe = pipe.to(device)
+
+        # pipe.enable_model_cpu_offload()
+
+        self.pipe = pipe
+        self.pipe.set_progress_bar_config(disable=True)
+
+        #     vae = AutoencoderKL.from_pretrained(
+        #         base_model_path, subfolder="vae", torch_dtype=torch_dtype
+        # )
+
+    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
+        generator = torch.manual_seed(params.seed)
+        steps = params.steps
+        strength = params.strength
+        prompt = params.prompt
+        guidance = params.guidance
+
+        results = self.pipe(
+            image=params.image,
+            prompt=prompt,
+            generator=generator,
+            strength=strength,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            width=params.width,
+            height=params.height,
+        )
+        return results.images[0]

server/requirements.txt

@@ -15,9 +15,16 @@ xformers; sys_platform != 'darwin' or platform_machine != 'arm64'
 markdown2
 safetensors
 stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/nightly/stable_fast-1.0.5.dev20241127+torch230cu121-cp310-cp310-manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
-oneflow @ https://github.com/siliconflow/oneflow_releases/releases/download/community_cu121/oneflow-0.9.1.dev20241114%2Bcu121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
-onediff @ git+https://github.com/siliconflow/onediff.git@main#egg=onediff ; sys_platform != 'darwin' or platform_machine != 'arm64'
+#oneflow @ https://github.com/siliconflow/oneflow_releases/releases/download/community_cu121/oneflow-0.9.1.dev20241114%2Bcu121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
+#onediff @ git+https://github.com/siliconflow/onediff.git@main#egg=onediff ; sys_platform != 'darwin' or platform_machine != 'arm64'
 setuptools
 mpmath==1.3.0
 numpy==1.*
-controlnet-aux
+controlnet-aux
+sentencepiece==0.2.0
+optimum-quanto
+gguf==0.13.0
+pydantic>=2.7.0
+types-Pillow
+mypy
+python-dotenv

server/util.py

@@ -1,10 +1,28 @@
 from importlib import import_module
-from types import ModuleType
+from typing import Any, TypeVar, Generic, TypeVar
 from PIL import Image
 import io
+from pydantic import BaseModel, create_model, Field
 
 
-def get_pipeline_class(pipeline_name: str) -> ModuleType:
+TPipeline = TypeVar("TPipeline", bound=type[Any])
+T = TypeVar('T')
+
+
+class ParamsModel(BaseModel):
+    """Base model for pipeline parameters."""
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> 'ParamsModel':
+        """Create a model instance from dictionary data."""
+        return cls.model_validate(data)
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert model to dictionary."""
+        return self.model_dump()
+
+
+def get_pipeline_class(pipeline_name: str) -> TPipeline:
     try:
         module = import_module(f"pipelines.{pipeline_name}")
     except ModuleNotFoundError:
@@ -15,6 +33,10 @@ def get_pipeline_class(pipeline_name: str) -> ModuleType:
     if pipeline_class is None:
         raise ValueError(f"'Pipeline' class not found in module '{pipeline_name}'.")
 
+    # Type check to ensure we're returning a class
+    if not isinstance(pipeline_class, type):
+        raise TypeError(f"'Pipeline' in module '{pipeline_name}' is not a class")
+
     return pipeline_class