new base

.gitignore

@@ -1,2 +1,4 @@
 __pycache__/
-venv/
+venv/
+public/
+*.pem

app.py

@@ -0,0 +1,14 @@
+from fastapi import FastAPI
+
+from config import args
+from device import device, torch_dtype
+from app_init import init_app
+from user_queue import user_queue_map
+from util import get_pipeline_class
+
+
+app = FastAPI()
+
+pipeline_class = get_pipeline_class(args.pipeline)
+pipeline = pipeline_class(args, device, torch_dtype)
+init_app(app, user_queue_map, args, pipeline)

app_init.py

@@ -0,0 +1,145 @@
+from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+import logging
+import traceback
+from config import Args
+from user_queue import UserQueueDict
+import uuid
+import asyncio
+import time
+from PIL import Image
+import io
+
+
+def init_app(app: FastAPI, user_queue_map: UserQueueDict, args: Args, pipeline):
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+    print("Init app", app)
+
+    @app.websocket("/ws")
+    async def websocket_endpoint(websocket: WebSocket):
+        await websocket.accept()
+        if args.max_queue_size > 0 and len(user_queue_map) >= args.max_queue_size:
+            print("Server is full")
+            await websocket.send_json({"status": "error", "message": "Server is full"})
+            await websocket.close()
+            return
+
+        try:
+            uid = uuid.uuid4()
+            print(f"New user connected: {uid}")
+            await websocket.send_json(
+                {"status": "success", "message": "Connected", "userId": uid}
+            )
+            user_queue_map[uid] = {"queue": asyncio.Queue()}
+            await websocket.send_json(
+                {"status": "start", "message": "Start Streaming", "userId": uid}
+            )
+            await handle_websocket_data(websocket, uid)
+        except WebSocketDisconnect as e:
+            logging.error(f"WebSocket Error: {e}, {uid}")
+            traceback.print_exc()
+        finally:
+            print(f"User disconnected: {uid}")
+            queue_value = user_queue_map.pop(uid, None)
+            queue = queue_value.get("queue", None)
+            if queue:
+                while not queue.empty():
+                    try:
+                        queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        continue
+
+    @app.get("/queue_size")
+    async def get_queue_size():
+        queue_size = len(user_queue_map)
+        return JSONResponse({"queue_size": queue_size})
+
+    @app.get("/stream/{user_id}")
+    async def stream(user_id: uuid.UUID):
+        uid = user_id
+        try:
+            user_queue = user_queue_map[uid]
+            queue = user_queue["queue"]
+
+            async def generate():
+                last_prompt: str = None
+                while True:
+                    data = await queue.get()
+                    input_image = data["image"]
+                    params = data["params"]
+                    if input_image is None:
+                        continue
+
+                    image = pipeline.predict(
+                        input_image,
+                        params,
+                    )
+                    if image is None:
+                        continue
+                    frame_data = io.BytesIO()
+                    image.save(frame_data, format="JPEG")
+                    frame_data = frame_data.getvalue()
+                    if frame_data is not None and len(frame_data) > 0:
+                        yield b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + frame_data + b"\r\n"
+
+                    await asyncio.sleep(1.0 / 120.0)
+
+            return StreamingResponse(
+                generate(), media_type="multipart/x-mixed-replace;boundary=frame"
+            )
+        except Exception as e:
+            logging.error(f"Streaming Error: {e}, {user_queue_map}")
+            traceback.print_exc()
+            return HTTPException(status_code=404, detail="User not found")
+
+    async def handle_websocket_data(websocket: WebSocket, user_id: uuid.UUID):
+        uid = user_id
+        user_queue = user_queue_map[uid]
+        queue = user_queue["queue"]
+        if not queue:
+            return HTTPException(status_code=404, detail="User not found")
+        last_time = time.time()
+        try:
+            while True:
+                data = await websocket.receive_bytes()
+                params = await websocket.receive_json()
+                params = pipeline.InputParams(**params)
+                pil_image = Image.open(io.BytesIO(data))
+
+                while not queue.empty():
+                    try:
+                        queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        continue
+                await queue.put({"image": pil_image, "params": params})
+                if args.timeout > 0 and time.time() - last_time > args.timeout:
+                    await websocket.send_json(
+                        {
+                            "status": "timeout",
+                            "message": "Your session has ended",
+                            "userId": uid,
+                        }
+                    )
+                    await websocket.close()
+                    return
+
+        except Exception as e:
+            logging.error(f"Error: {e}")
+            traceback.print_exc()
+
+    # route to setup frontend
+    @app.get("/settings")
+    async def settings():
+        params = pipeline.InputParams()
+        return JSONResponse({"settings": params.dict()})
+
+    app.mount("/", StaticFiles(directory="public", html=True), name="public")

build-run.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+cd frontend
+npm install
+npm run build
+if [ $? -eq 0 ]; then
+    echo -e "\033[1;32m\nfrontend build success \033[0m"
+else
+    echo -e "\033[1;31m\nfrontend build failed\n\033[0m" >&2  exit 1
+fi
+cd ../
+python run.py --reload
+

config.py

@@ -0,0 +1,58 @@
+from typing import NamedTuple
+import argparse
+import os
+
+
+class Args(NamedTuple):
+    host: str
+    port: int
+    reload: bool
+    mode: str
+    max_queue_size: int
+    timeout: float
+    safety_checker: bool
+    torch_compile: bool
+    use_taesd: bool
+    pipeline: str
+
+
+MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
+TIMEOUT = float(os.environ.get("TIMEOUT", 0))
+SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None) == "True"
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None) == "True"
+USE_TAESD = os.environ.get("USE_TAESD", None) == "True"
+default_host = os.getenv("HOST", "0.0.0.0")
+default_port = int(os.getenv("PORT", "7860"))
+default_mode = os.getenv("MODE", "default")
+
+parser = argparse.ArgumentParser(description="Run the app")
+parser.add_argument("--host", type=str, default=default_host, help="Host address")
+parser.add_argument("--port", type=int, default=default_port, help="Port number")
+parser.add_argument("--reload", action="store_true", help="Reload code on change")
+parser.add_argument(
+    "--mode", type=str, default=default_mode, help="App Inferece Mode: txt2img, img2img"
+)
+parser.add_argument(
+    "--max_queue_size", type=int, default=MAX_QUEUE_SIZE, help="Max Queue Size"
+)
+parser.add_argument("--timeout", type=float, default=TIMEOUT, help="Timeout")
+parser.add_argument(
+    "--safety_checker", type=bool, default=SAFETY_CHECKER, help="Safety Checker"
+)
+parser.add_argument(
+    "--torch_compile", type=bool, default=TORCH_COMPILE, help="Torch Compile"
+)
+parser.add_argument(
+    "--use_taesd",
+    type=bool,
+    default=USE_TAESD,
+    help="Use Tiny Autoencoder",
+)
+parser.add_argument(
+    "--pipeline",
+    type=str,
+    default="txt2img",
+    help="Pipeline to use",
+)
+
+args = Args(**vars(parser.parse_args()))

device.py

@@ -0,0 +1,12 @@
+import torch
+
+# check if MPS is available OSX only M1/M2/M3 chips
+mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
+device = torch.device(
+    "cuda" if torch.cuda.is_available() else "xpu" if xpu_available else "cpu"
+)
+torch_dtype = torch.float16
+if mps_available:
+    device = torch.device("mps")
+    torch_dtype = torch.float32

frontend/.eslintignore

@@ -0,0 +1,13 @@
+.DS_Store
+node_modules
+/build
+/.svelte-kit
+/package
+.env
+.env.*
+!.env.example
+
+# Ignore files for PNPM, NPM and YARN
+pnpm-lock.yaml
+package-lock.json
+yarn.lock

frontend/.eslintrc.cjs

@@ -0,0 +1,30 @@
+module.exports = {
+  root: true,
+  extends: [
+    'eslint:recommended',
+    'plugin:@typescript-eslint/recommended',
+    'plugin:svelte/recommended',
+    'prettier'
+  ],
+  parser: '@typescript-eslint/parser',
+  plugins: ['@typescript-eslint'],
+  parserOptions: {
+    sourceType: 'module',
+    ecmaVersion: 2020,
+    extraFileExtensions: ['.svelte']
+  },
+  env: {
+    browser: true,
+    es2017: true,
+    node: true
+  },
+  overrides: [
+    {
+      files: ['*.svelte'],
+      parser: 'svelte-eslint-parser',
+      parserOptions: {
+        parser: '@typescript-eslint/parser'
+      }
+    }
+  ]
+};

frontend/.gitignore

@@ -0,0 +1,10 @@
+.DS_Store
+node_modules
+/build
+/.svelte-kit
+/package
+.env
+.env.*
+!.env.example
+vite.config.js.timestamp-*
+vite.config.ts.timestamp-*

frontend/.npmrc

@@ -0,0 +1 @@
+engine-strict=true

frontend/.prettierignore

@@ -0,0 +1,13 @@
+.DS_Store
+node_modules
+/build
+/.svelte-kit
+/package
+.env
+.env.*
+!.env.example
+
+# Ignore files for PNPM, NPM and YARN
+pnpm-lock.yaml
+package-lock.json
+yarn.lock

frontend/.prettierrc

@@ -0,0 +1,19 @@
+{
+  "useTabs": false,
+  "singleQuote": true,
+  "trailingComma": "none",
+  "printWidth": 100,
+  "plugins": [
+    "prettier-plugin-svelte",
+    "prettier-plugin-organize-imports",
+    "prettier-plugin-tailwindcss"
+  ],
+  "overrides": [
+    {
+      "files": "*.svelte",
+      "options": {
+        "parser": "svelte"
+      }
+    }
+  ]
+}

frontend/README.md

@@ -0,0 +1,38 @@
+# create-svelte
+
+Everything you need to build a Svelte project, powered by [`create-svelte`](https://github.com/sveltejs/kit/tree/master/packages/create-svelte).
+
+## Creating a project
+
+If you're seeing this, you've probably already done this step. Congrats!
+
+```bash
+# create a new project in the current directory
+npm create svelte@latest
+
+# create a new project in my-app
+npm create svelte@latest my-app
+```
+
+## Developing
+
+Once you've created a project and installed dependencies with `npm install` (or `pnpm install` or `yarn`), start a development server:
+
+```bash
+npm run dev
+
+# or start the server and open the app in a new browser tab
+npm run dev -- --open
+```
+
+## Building
+
+To create a production version of your app:
+
+```bash
+npm run build
+```
+
+You can preview the production build with `npm run preview`.
+
+> To deploy your app, you may need to install an [adapter](https://kit.svelte.dev/docs/adapters) for your target environment.

frontend/package.json

@@ -0,0 +1,36 @@
+{
+  "name": "frontend",
+  "version": "0.0.1",
+  "private": true,
+  "scripts": {
+    "dev": "vite dev",
+    "build": "vite build",
+    "preview": "vite preview",
+    "check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json",
+    "check:watch": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json --watch",
+    "lint": "prettier --check . && eslint .",
+    "format": "prettier --write ."
+  },
+  "devDependencies": {
+    "@sveltejs/adapter-auto": "^2.0.0",
+    "@sveltejs/kit": "^1.20.4",
+    "@typescript-eslint/eslint-plugin": "^6.0.0",
+    "@typescript-eslint/parser": "^6.0.0",
+    "autoprefixer": "^10.4.16",
+    "eslint": "^8.28.0",
+    "eslint-config-prettier": "^9.0.0",
+    "eslint-plugin-svelte": "^2.30.0",
+    "postcss": "^8.4.31",
+    "prettier": "^3.1.0",
+    "prettier-plugin-organize-imports": "^3.2.4",
+    "prettier-plugin-svelte": "^3.1.0",
+    "prettier-plugin-tailwindcss": "^0.5.7",
+    "svelte": "^4.0.5",
+    "svelte-check": "^3.4.3",
+    "tailwindcss": "^3.3.5",
+    "tslib": "^2.4.1",
+    "typescript": "^5.0.0",
+    "vite": "^4.4.2"
+  },
+  "type": "module"
+}

frontend/postcss.config.js

@@ -0,0 +1,6 @@
+export default {
+  plugins: {
+    tailwindcss: {},
+    autoprefixer: {}
+  }
+};

frontend/src/app.css

@@ -0,0 +1,3 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;

frontend/src/app.d.ts

@@ -0,0 +1,12 @@
+// See https://kit.svelte.dev/docs/types#app
+// for information about these interfaces
+declare global {
+  namespace App {
+    // interface Error {}
+    // interface Locals {}
+    // interface PageData {}
+    // interface Platform {}
+  }
+}
+
+export {};

frontend/src/app.html

@@ -0,0 +1,12 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <link rel="icon" href="%sveltekit.assets%/favicon.png" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    %sveltekit.head%
+  </head>
+  <body data-sveltekit-preload-data="hover">
+    <div style="display: contents">%sveltekit.body%</div>
+  </body>
+</html>

frontend/src/lib/index.ts

@@ -0,0 +1 @@
+// place files you want to import through the `$lib` alias in this folder.

frontend/src/routes/+layout.svelte

@@ -0,0 +1,5 @@
+<script>
+  import '../app.css';
+</script>
+
+<slot />

frontend/src/routes/+page.svelte

@@ -0,0 +1,160 @@
+<script lang="ts">
+  import { onMount } from 'svelte';
+  import { PUBLIC_BASE_URL } from '$env/static/public';
+
+  onMount(() => {
+    getSettings();
+  });
+  async function getSettings() {
+    const settings = await fetch(`${PUBLIC_BASE_URL}/settings`).then((r) => r.json());
+    console.log(settings);
+  }
+</script>
+
+<div class="fixed right-2 top-2 max-w-xs rounded-lg p-4 text-center text-sm font-bold" id="error" />
+<main class="container mx-auto flex max-w-4xl flex-col gap-4 px-4 py-4">
+  <article class="mx-auto max-w-xl text-center">
+    <h1 class="text-3xl font-bold">Real-Time Latent Consistency Model</h1>
+    <h2 class="mb-4 text-2xl font-bold">Image to Image</h2>
+    <p class="text-sm">
+      This demo showcases
+      <a
+        href="https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7"
+        target="_blank"
+        class="text-blue-500 underline hover:no-underline">LCM</a
+      >
+      Image to Image pipeline using
+      <a
+        href="https://github.com/huggingface/diffusers/tree/main/examples/community#latent-consistency-pipeline"
+        target="_blank"
+        class="text-blue-500 underline hover:no-underline">Diffusers</a
+      > with a MJPEG stream server.
+    </p>
+    <p class="text-sm">
+      There are <span id="queue_size" class="font-bold">0</span> user(s) sharing the same GPU,
+      affecting real-time performance. Maximum queue size is 4.
+      <a
+        href="https://huggingface.co/spaces/radames/Real-Time-Latent-Consistency-Model?duplicate=true"
+        target="_blank"
+        class="text-blue-500 underline hover:no-underline">Duplicate</a
+      > and run it on your own GPU.
+    </p>
+  </article>
+  <div>
+    <h2 class="font-medium">Prompt</h2>
+    <p class="text-sm text-gray-500">
+      Change the prompt to generate different images, accepts <a
+        href="https://github.com/damian0815/compel/blob/main/doc/syntax.md"
+        target="_blank"
+        class="text-blue-500 underline hover:no-underline">Compel</a
+      > syntax.
+    </p>
+    <div class="text-normal flex items-center rounded-md border border-gray-700 px-1 py-1">
+      <textarea
+        type="text"
+        id="prompt"
+        class="mx-1 w-full px-3 py-2 font-light outline-none dark:text-black"
+        title="Prompt, this is an example, feel free to modify"
+        placeholder="Add your prompt here..."
+        >Portrait of The Terminator with , glare pose, detailed, intricate, full of colour,
+        cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details,
+        unreal engine 5, cinematic, masterpiece</textarea
+      >
+    </div>
+  </div>
+  <div class="">
+    <details>
+      <summary class="cursor-pointer font-medium">Advanced Options</summary>
+      <div class="grid max-w-md grid-cols-3 items-center gap-3 py-3">
+        <label class="text-sm font-medium" for="guidance-scale">Guidance Scale </label>
+        <input
+          type="range"
+          id="guidance-scale"
+          name="guidance-scale"
+          min="1"
+          max="30"
+          step="0.001"
+          value="8.0"
+          oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)"
+        />
+        <output
+          class="w-[50px] rounded-md border border-gray-700 px-1 py-1 text-center text-xs font-light"
+        >
+          8.0</output
+        >
+        <label class="text-sm font-medium" for="strength">Strength</label>
+        <input
+          type="range"
+          id="strength"
+          name="strength"
+          min="0.20"
+          max="1"
+          step="0.001"
+          value="0.50"
+          oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)"
+        />
+        <output
+          class="w-[50px] rounded-md border border-gray-700 px-1 py-1 text-center text-xs font-light"
+        >
+          0.5</output
+        >
+        <label class="text-sm font-medium" for="seed">Seed</label>
+        <input
+          type="number"
+          id="seed"
+          name="seed"
+          value="299792458"
+          class="rounded-md border border-gray-700 p-2 text-right font-light dark:text-black"
+        />
+        <button
+          onclick="document.querySelector('#seed').value = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER)"
+          class="button"
+        >
+          Rand
+        </button>
+      </div>
+    </details>
+  </div>
+  <div class="flex gap-3">
+    <button id="start" class="button"> Start </button>
+    <button id="stop" class="button"> Stop </button>
+    <button id="snap" disabled class="button ml-auto"> Snapshot </button>
+  </div>
+  <div class="relative overflow-hidden rounded-lg border border-slate-300">
+    <img
+      id="player"
+      class="aspect-square w-full rounded-lg"
+      src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
+    />
+    <div class="absolute left-0 top-0 aspect-square w-1/4">
+      <video
+        id="webcam"
+        class="relative z-10 aspect-square w-full object-cover"
+        playsinline
+        autoplay
+        muted
+        loop
+      />
+      <svg
+        xmlns="http://www.w3.org/2000/svg"
+        viewBox="0 0 448 448"
+        width="100"
+        class="absolute top-0 z-0 w-full p-4 opacity-20"
+      >
+        <path
+          fill="currentColor"
+          d="M224 256a128 128 0 1 0 0-256 128 128 0 1 0 0 256zm-45.7 48A178.3 178.3 0 0 0 0 482.3 29.7 29.7 0 0 0 29.7 512h388.6a29.7 29.7 0 0 0 29.7-29.7c0-98.5-79.8-178.3-178.3-178.3h-91.4z"
+        />
+      </svg>
+    </div>
+  </div>
+</main>
+
+<style lang="postcss">
+  :global(html) {
+    @apply text-black dark:bg-gray-900 dark:text-white;
+  }
+  .button {
+    @apply rounded bg-gray-700 p-2 font-normal text-white hover:bg-gray-800 disabled:cursor-not-allowed disabled:bg-gray-300 dark:disabled:bg-gray-700 dark:disabled:text-black;
+  }
+</style>

frontend/src/routes/+page.ts

@@ -0,0 +1 @@
+export const prerender = true

frontend/svelte.config.js

@@ -0,0 +1,19 @@
+import adapter from '@sveltejs/adapter-static';
+import { vitePreprocess } from '@sveltejs/kit/vite';
+
+/** @type {import('@sveltejs/kit').Config} */
+const config = {
+  preprocess: vitePreprocess(),
+
+  kit: {
+    adapter: adapter({
+      pages: '../public',
+      assets: '../public',
+      fallback: undefined,
+      precompress: false,
+      strict: true
+    })
+  }
+};
+
+export default config;

frontend/tailwind.config.js

@@ -0,0 +1,8 @@
+/** @type {import('tailwindcss').Config} */
+export default {
+  content: ['./src/**/*.{html,js,svelte,ts}'],
+  theme: {
+    extend: {}
+  },
+  plugins: []
+};

frontend/tsconfig.json

@@ -0,0 +1,17 @@
+{
+  "extends": "./.svelte-kit/tsconfig.json",
+  "compilerOptions": {
+    "allowJs": true,
+    "checkJs": true,
+    "esModuleInterop": true,
+    "forceConsistentCasingInFileNames": true,
+    "resolveJsonModule": true,
+    "skipLibCheck": true,
+    "sourceMap": true,
+    "strict": true
+  }
+  // Path aliases are handled by https://kit.svelte.dev/docs/configuration#alias
+  //
+  // If you want to overwrite includes/excludes, make sure to copy over the relevant includes/excludes
+  // from the referenced tsconfig.json - TypeScript does not merge them in
+}

frontend/vite.config.ts

@@ -0,0 +1,6 @@
+import { sveltekit } from '@sveltejs/kit/vite';
+import { defineConfig } from 'vite';
+
+export default defineConfig({
+  plugins: [sveltekit()]
+});

pipelines/controlnet.py

@@ -0,0 +1,90 @@
+from diffusers import DiffusionPipeline, AutoencoderTiny
+from latent_consistency_controlnet import LatentConsistencyModelPipeline_controlnet
+
+from compel import Compel
+import torch
+
+try:
+    import intel_extension_for_pytorch as ipex  # type: ignore
+except:
+    pass
+
+import psutil
+from config import Args
+from pydantic import BaseModel
+from PIL import Image
+from typing import Callable
+
+base_model = "SimianLuo/LCM_Dreamshaper_v7"
+WIDTH = 512
+HEIGHT = 512
+
+
+class Pipeline:
+    class InputParams(BaseModel):
+        seed: int = 2159232
+        prompt: str
+        guidance_scale: float = 8.0
+        strength: float = 0.5
+        steps: int = 4
+        lcm_steps: int = 50
+        width: int = WIDTH
+        height: int = HEIGHT
+
+    @staticmethod
+    def create_pipeline(
+        args: Args, device: torch.device, torch_dtype: torch.dtype
+    ) -> Callable[["Pipeline.InputParams"], Image.Image]:
+        if args.safety_checker:
+            pipe = DiffusionPipeline.from_pretrained(base_model)
+        else:
+            pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None)
+        if args.use_taesd:
+            pipe.vae = AutoencoderTiny.from_pretrained(
+                "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True
+            )
+
+        pipe.set_progress_bar_config(disable=True)
+        pipe.to(device=device, dtype=torch_dtype)
+        pipe.unet.to(memory_format=torch.channels_last)
+
+        # check if computer has less than 64GB of RAM using sys or os
+        if psutil.virtual_memory().total < 64 * 1024**3:
+            pipe.enable_attention_slicing()
+
+        if args.torch_compile:
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
+
+            pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+
+        compel_proc = Compel(
+            tokenizer=pipe.tokenizer,
+            text_encoder=pipe.text_encoder,
+            truncate_long_prompts=False,
+        )
+
+        def predict(params: "Pipeline.InputParams") -> Image.Image:
+            generator = torch.manual_seed(params.seed)
+            prompt_embeds = compel_proc(params.prompt)
+            # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
+            results = pipe(
+                prompt_embeds=prompt_embeds,
+                generator=generator,
+                num_inference_steps=params.steps,
+                guidance_scale=params.guidance_scale,
+                width=params.width,
+                height=params.height,
+                original_inference_steps=params.lcm_steps,
+                output_type="pil",
+            )
+            nsfw_content_detected = (
+                results.nsfw_content_detected[0]
+                if "nsfw_content_detected" in results
+                else False
+            )
+            if nsfw_content_detected:
+                return None
+            return results.images[0]
+
+        return predict

pipelines/txt2img.py

@@ -0,0 +1,85 @@
+from diffusers import DiffusionPipeline, AutoencoderTiny
+from compel import Compel
+import torch
+
+try:
+    import intel_extension_for_pytorch as ipex  # type: ignore
+except:
+    pass
+
+import psutil
+from config import Args
+from pydantic import BaseModel
+from PIL import Image
+from typing import Callable
+
+base_model = "SimianLuo/LCM_Dreamshaper_v7"
+taesd_model = "madebyollin/taesd"
+
+
+class Pipeline:
+    class InputParams(BaseModel):
+        seed: int = 2159232
+        prompt: str = ""
+        guidance_scale: float = 8.0
+        strength: float = 0.5
+        steps: int = 4
+        width: int = 512
+        height: int = 512
+
+    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
+        if args.safety_checker:
+            self.pipe = DiffusionPipeline.from_pretrained(base_model)
+        else:
+            self.pipe = DiffusionPipeline.from_pretrained(
+                base_model, safety_checker=None
+            )
+        if args.use_taesd:
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+            )
+
+        self.pipe.set_progress_bar_config(disable=True)
+        self.pipe.to(device=device, dtype=torch_dtype)
+        self.pipe.unet.to(memory_format=torch.channels_last)
+
+        # check if computer has less than 64GB of RAM using sys or os
+        if psutil.virtual_memory().total < 64 * 1024**3:
+            self.pipe.enable_attention_slicing()
+
+        if args.torch_compile:
+            self.pipe.unet = torch.compile(
+                self.pipe.unet, mode="reduce-overhead", fullgraph=True
+            )
+            self.pipe.vae = torch.compile(
+                self.pipe.vae, mode="reduce-overhead", fullgraph=True
+            )
+
+            self.pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+
+        self.compel_proc = Compel(
+            tokenizer=self.pipe.tokenizer,
+            text_encoder=self.pipe.text_encoder,
+            truncate_long_prompts=False,
+        )
+
+    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
+        generator = torch.manual_seed(params.seed)
+        prompt_embeds = self.compel_proc(params.prompt)
+        results = self.pipe(
+            prompt_embeds=prompt_embeds,
+            generator=generator,
+            num_inference_steps=params.steps,
+            guidance_scale=params.guidance_scale,
+            width=params.width,
+            height=params.height,
+            output_type="pil",
+        )
+        nsfw_content_detected = (
+            results.nsfw_content_detected[0]
+            if "nsfw_content_detected" in results
+            else False
+        )
+        if nsfw_content_detected:
+            return None
+        return results.images[0]

pipelines/txt2imglora.py

@@ -0,0 +1,93 @@
+from diffusers import DiffusionPipeline, AutoencoderTiny
+from compel import Compel
+import torch
+
+try:
+    import intel_extension_for_pytorch as ipex  # type: ignore
+except:
+    pass
+
+import psutil
+from config import Args
+from pydantic import BaseModel
+from PIL import Image
+from typing import Callable
+
+base_model = "SimianLuo/LCM_Dreamshaper_v7"
+WIDTH = 512
+HEIGHT = 512
+
+model_id = "wavymulder/Analog-Diffusion"
+lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
+
+
+class Pipeline:
+    class InputParams(BaseModel):
+        seed: int = 2159232
+        prompt: str
+        guidance_scale: float = 8.0
+        strength: float = 0.5
+        steps: int = 4
+        lcm_steps: int = 50
+        width: int = WIDTH
+        height: int = HEIGHT
+
+    @staticmethod
+    def create_pipeline(
+        args: Args, device: torch.device, torch_dtype: torch.dtype
+    ) -> Callable[["Pipeline.InputParams"], Image.Image]:
+        if args.safety_checker:
+            pipe = DiffusionPipeline.from_pretrained(base_model)
+        else:
+            pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None)
+        if args.use_taesd:
+            pipe.vae = AutoencoderTiny.from_pretrained(
+                "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True
+            )
+
+        pipe.set_progress_bar_config(disable=True)
+        pipe.to(device=device, dtype=torch_dtype)
+        pipe.unet.to(memory_format=torch.channels_last)
+
+        # Load LCM LoRA
+        pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm")
+        # check if computer has less than 64GB of RAM using sys or os
+        if psutil.virtual_memory().total < 64 * 1024**3:
+            pipe.enable_attention_slicing()
+
+        if args.torch_compile:
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
+
+            pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+
+        compel_proc = Compel(
+            tokenizer=pipe.tokenizer,
+            text_encoder=pipe.text_encoder,
+            truncate_long_prompts=False,
+        )
+
+        def predict(params: "Pipeline.InputParams") -> Image.Image:
+            generator = torch.manual_seed(params.seed)
+            prompt_embeds = compel_proc(params.prompt)
+            # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
+            results = pipe(
+                prompt_embeds=prompt_embeds,
+                generator=generator,
+                num_inference_steps=params.steps,
+                guidance_scale=params.guidance_scale,
+                width=params.width,
+                height=params.height,
+                original_inference_steps=params.lcm_steps,
+                output_type="pil",
+            )
+            nsfw_content_detected = (
+                results.nsfw_content_detected[0]
+                if "nsfw_content_detected" in results
+                else False
+            )
+            if nsfw_content_detected:
+                return None
+            return results.images[0]
+
+        return predict

requirements.txt

@@ -3,8 +3,8 @@ transformers==4.34.1
 gradio==3.50.2
 --extra-index-url https://download.pytorch.org/whl/cu121;
 torch==2.1.0
-fastapi==0.104.0
-uvicorn==0.23.2
+fastapi==0.104.1
+uvicorn==0.24.0.post1
 Pillow==10.1.0
 accelerate==0.24.0
 compel==2.0.2

run.py

@@ -0,0 +1,5 @@
+if __name__ == "__main__":
+    import uvicorn
+    from config import args
+
+    uvicorn.run("app:app", host=args.host, port=args.port, reload=args.reload)

user_queue.py

@@ -0,0 +1,18 @@
+from typing import Dict, Union
+from uuid import UUID
+from asyncio import Queue
+from PIL import Image
+from typing import Tuple, Union
+from uuid import UUID
+from asyncio import Queue
+from PIL import Image
+
+UserId = UUID
+
+InputParams = dict
+
+QueueContent = Dict[str, Union[Image.Image, InputParams]]
+
+UserQueueDict = Dict[UserId, Queue[QueueContent]]
+
+user_queue_map: UserQueueDict = {}

util.py

@@ -0,0 +1,16 @@
+from importlib import import_module
+from types import ModuleType
+
+
+def get_pipeline_class(pipeline_name: str) -> ModuleType:
+    try:
+        module = import_module(f"pipelines.{pipeline_name}")
+    except ModuleNotFoundError:
+        raise ValueError(f"Pipeline {pipeline_name} module not found")
+
+    pipeline_class = getattr(module, "Pipeline", None)
+
+    if pipeline_class is None:
+        raise ValueError(f"'Pipeline' class not found in module '{module_name}'.")
+
+    return pipeline_class