add contronet canny

app-controlnet.py

@@ -0,0 +1,296 @@
+import asyncio
+import json
+import logging
+import traceback
+from pydantic import BaseModel
+
+from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+
+from diffusers import AutoencoderTiny, ControlNetModel
+from latent_consistency_controlnet import LatentConsistencyModelPipeline_controlnet
+from compel import Compel
+import torch
+
+from canny_gpu import SobelOperator 
+# from controlnet_aux import OpenposeDetector
+# import cv2
+
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
+from PIL import Image
+import numpy as np
+import gradio as gr
+import io
+import uuid
+import os
+import time
+import psutil
+
+
+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)
+WIDTH = 512
+HEIGHT = 512
+# disable tiny autoencoder for better quality speed tradeoff
+USE_TINY_AUTOENCODER = True
+
+# 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"
+)
+
+# change to torch.float16 to save GPU memory
+torch_dtype = torch.float16
+
+print(f"TIMEOUT: {TIMEOUT}")
+print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
+print(f"MAX_QUEUE_SIZE: {MAX_QUEUE_SIZE}")
+print(f"device: {device}")
+
+if mps_available:
+    device = torch.device("mps")
+    device = "cpu"
+    torch_dtype = torch.float32
+
+controlnet_canny = ControlNetModel.from_pretrained(
+    "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch_dtype
+).to(device)
+
+canny_torch = SobelOperator(device=device)
+# controlnet_pose = ControlNetModel.from_pretrained(
+#     "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch_dtype
+# ).to(device)
+# controlnet_depth = ControlNetModel.from_pretrained(
+#     "lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch_dtype
+# ).to(device)
+
+
+# pose_processor = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+
+if SAFETY_CHECKER == "True":
+    pipe = LatentConsistencyModelPipeline_controlnet.from_pretrained(
+        "SimianLuo/LCM_Dreamshaper_v7",
+        controlnet=controlnet_canny,
+        scheduler=None,
+    )
+else:
+    pipe = LatentConsistencyModelPipeline_controlnet.from_pretrained(
+        "SimianLuo/LCM_Dreamshaper_v7",
+        safety_checker=None,
+        controlnet=controlnet_canny,
+        scheduler=None,
+    )
+
+if USE_TINY_AUTOENCODER:
+    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).to(device)
+pipe.unet.to(memory_format=torch.channels_last)
+
+if psutil.virtual_memory().total < 64 * 1024**3:
+    pipe.enable_attention_slicing()
+
+# if not mps_available and not xpu_available:
+#     pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+#     pipe(prompt="warmup", image=[Image.new("RGB", (512, 512))])
+
+compel_proc = Compel(
+    tokenizer=pipe.tokenizer,
+    text_encoder=pipe.text_encoder,
+    truncate_long_prompts=False,
+)
+user_queue_map = {}
+
+
+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
+    controlnet_scale: float = 0.8
+    controlnet_start: float = 0.0
+    controlnet_end: float = 1.0
+    canny_low_threshold: float = 0.31
+    canny_high_threshold: float = 0.78
+
+def predict(
+    input_image: Image.Image, params: InputParams, prompt_embeds: torch.Tensor = None
+):
+    generator = torch.manual_seed(params.seed)
+    
+    control_image = canny_torch(input_image, params.canny_low_threshold, params.canny_high_threshold)
+    print(params.canny_low_threshold, params.canny_high_threshold)
+    results = pipe(
+        control_image=control_image,
+        prompt_embeds=prompt_embeds,
+        generator=generator,
+        image=input_image,
+        strength=params.strength,
+        num_inference_steps=params.steps,
+        guidance_scale=params.guidance_scale,
+        width=params.width,
+        height=params.height,
+        lcm_origin_steps=params.lcm_steps,
+        output_type="pil",
+        controlnet_conditioning_scale=params.controlnet_scale,
+        control_guidance_start=params.controlnet_start,
+        control_guidance_end=params.controlnet_end,
+    )
+    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]
+
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    if MAX_QUEUE_SIZE > 0 and len(user_queue_map) >= MAX_QUEUE_SIZE:
+        print("Server is full")
+        await websocket.send_json({"status": "error", "message": "Server is full"})
+        await websocket.close()
+        return
+
+    try:
+        uid = str(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 = str(user_id)
+    try:
+        user_queue = user_queue_map[uid]
+        queue = user_queue["queue"]
+
+        async def generate():
+            last_prompt: str = None
+            prompt_embeds: torch.Tensor = None
+            while True:
+                data = await queue.get()
+                input_image = data["image"]
+                params = data["params"]
+                if input_image is None:
+                    continue
+                # avoid recalculate prompt embeds
+                if last_prompt != params.prompt:
+                    print("new prompt")
+                    prompt_embeds = compel_proc(params.prompt)
+                    last_prompt = params.prompt
+
+                image = predict(
+                    input_image,
+                    params,
+                    prompt_embeds,
+                )
+                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 = str(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 = 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 TIMEOUT > 0 and time.time() - last_time > 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()
+
+
+app.mount("/", StaticFiles(directory="controlnet", html=True), name="public")

canny_gpu.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+from torchvision.transforms import ToTensor, ToPILImage
+from PIL import Image
+
+class SobelOperator(nn.Module):
+    def __init__(self, device="cuda"):
+        super(SobelOperator, self).__init__()
+        self.device = device
+        self.edge_conv_x = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to(
+            self.device
+        )
+        self.edge_conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to(
+            self.device
+        )
+
+        sobel_kernel_x = torch.tensor(
+            [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=self.device
+        )
+        sobel_kernel_y = torch.tensor(
+            [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]], device=self.device
+        )
+
+        self.edge_conv_x.weight = nn.Parameter(sobel_kernel_x.view((1, 1, 3, 3)))
+        self.edge_conv_y.weight = nn.Parameter(sobel_kernel_y.view((1, 1, 3, 3)))
+
+    @torch.no_grad()
+    def forward(self, image: Image.Image, low_threshold: float, high_threshold: float):
+        # Convert PIL image to PyTorch tensor
+        image_gray = image.convert("L")
+        image_tensor = ToTensor()(image_gray).unsqueeze(0).to(self.device)
+
+        # Compute gradients
+        edge_x = self.edge_conv_x(image_tensor)
+        edge_y = self.edge_conv_y(image_tensor)
+        edge = torch.sqrt(edge_x**2 + edge_y**2)
+
+        # Apply thresholding
+        edge = edge / edge.max()  # Normalize to 0-1
+        edge[edge >= high_threshold] = 1.0
+        edge[edge <= low_threshold] = 0.0
+
+        # Convert the result back to a PIL image
+        return ToPILImage()(edge.squeeze(0).cpu())

controlnet/index.html

@@ -0,0 +1,414 @@
+<!doctype html>
+<html>
+
+<head>
+    <meta charset="UTF-8">
+    <title>Real-Time Latent Consistency Model ControlNet</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <script
+        src="https://cdnjs.cloudflare.com/ajax/libs/iframe-resizer/4.3.1/iframeResizer.contentWindow.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/[email protected]/piexif.min.js"></script>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <style type="text/tailwindcss">
+        .button {
+          @apply bg-gray-700 hover:bg-gray-800 text-white font-normal p-2 rounded disabled:bg-gray-300 dark:disabled:bg-gray-700 disabled:cursor-not-allowed dark:disabled:text-black
+        }
+    </style>
+    <script type="module">
+        // you can change the size of the input image to 768x768 if you have a powerful GPU
+        const getValue = (id) => document.querySelector(`${id}`).value;
+        const startBtn = document.querySelector("#start");
+        const stopBtn = document.querySelector("#stop");
+        const videoEl = document.querySelector("#webcam");
+        const imageEl = document.querySelector("#player");
+        const queueSizeEl = document.querySelector("#queue_size");
+        const errorEl = document.querySelector("#error");
+        const snapBtn = document.querySelector("#snap");
+        const webcamsEl = document.querySelector("#webcams");
+
+        function LCMLive(webcamVideo, liveImage) {
+            let websocket;
+
+            async function start() {
+                return new Promise((resolve, reject) => {
+                    const websocketURL = `${window.location.protocol === "https:" ? "wss" : "ws"
+                        }:${window.location.host}/ws`;
+
+                    const socket = new WebSocket(websocketURL);
+                    socket.onopen = () => {
+                        console.log("Connected to websocket");
+                    };
+                    socket.onclose = () => {
+                        console.log("Disconnected from websocket");
+                        stop();
+                        resolve({ "status": "disconnected" });
+                    };
+                    socket.onerror = (err) => {
+                        console.error(err);
+                        reject(err);
+                    };
+                    socket.onmessage = (event) => {
+                        const data = JSON.parse(event.data);
+                        switch (data.status) {
+                            case "success":
+                                break;
+                            case "start":
+                                const userId = data.userId;
+                                initVideoStream(userId);
+                                break;
+                            case "timeout":
+                                stop();
+                                resolve({ "status": "timeout" });
+                            case "error":
+                                stop();
+                                reject(data.message);
+
+                        }
+                    };
+                    websocket = socket;
+                })
+            }
+            function switchCamera() {
+                const constraints = {
+                    audio: false,
+                    video: { width: 1024, height: 768, deviceId: mediaDevices[webcamsEl.value].deviceId }
+                };
+                navigator.mediaDevices
+                    .getUserMedia(constraints)
+                    .then((mediaStream) => {
+                        webcamVideo.removeEventListener("timeupdate", videoTimeUpdateHandler);
+                        webcamVideo.srcObject = mediaStream;
+                        webcamVideo.onloadedmetadata = () => {
+                            webcamVideo.play();
+                            webcamVideo.addEventListener("timeupdate", videoTimeUpdateHandler);
+                        };
+                    })
+                    .catch((err) => {
+                        console.error(`${err.name}: ${err.message}`);
+                    });
+            }
+
+            async function videoTimeUpdateHandler() {
+                const dimension = getValue("input[name=dimension]:checked");
+                const [WIDTH, HEIGHT] = JSON.parse(dimension);
+
+                const canvas = new OffscreenCanvas(WIDTH, HEIGHT);
+                const videoW = webcamVideo.videoWidth;
+                const videoH = webcamVideo.videoHeight;
+                const aspectRatio = WIDTH / HEIGHT;
+
+                const ctx = canvas.getContext("2d");
+                ctx.drawImage(webcamVideo, videoW / 2 - videoH * aspectRatio / 2, 0, videoH * aspectRatio, videoH, 0, 0, WIDTH, HEIGHT)
+                const blob = await canvas.convertToBlob({ type: "image/jpeg", quality: 1 });
+                websocket.send(blob);
+                websocket.send(JSON.stringify({
+                    "seed": getValue("#seed"),
+                    "prompt": getValue("#prompt"),
+                    "guidance_scale": getValue("#guidance-scale"),
+                    "strength": getValue("#strength"),
+                    "steps": getValue("#steps"),
+                    "lcm_steps": getValue("#lcm_steps"),
+                    "width": WIDTH,
+                    "height": HEIGHT,
+                    "controlnet_scale": getValue("#controlnet_scale"),
+                    "controlnet_start": getValue("#controlnet_start"),
+                    "controlnet_end": getValue("#controlnet_end"),
+                    "canny_low_threshold": getValue("#canny_low_threshold"),
+                    "canny_high_threshold": getValue("#canny_high_threshold"),
+                }));
+            }
+            let mediaDevices = [];
+            async function initVideoStream(userId) {
+                liveImage.src = `/stream/${userId}`;
+                await navigator.mediaDevices.enumerateDevices()
+                    .then(devices => {
+                        const cameras = devices.filter(device => device.kind === 'videoinput');
+                        mediaDevices = cameras;
+                        webcamsEl.innerHTML = "";
+                        cameras.forEach((camera, index) => {
+                            const option = document.createElement("option");
+                            option.value = index;
+                            option.innerText = camera.label;
+                            webcamsEl.appendChild(option);
+                            option.selected = index === 0;
+                        });
+                        webcamsEl.addEventListener("change", switchCamera);
+                    })
+                    .catch(err => {
+                        console.error(err);
+                    });
+                const constraints = {
+                    audio: false,
+                    video: { width: 1024, height: 768, deviceId: mediaDevices[0].deviceId }
+                };
+                navigator.mediaDevices
+                    .getUserMedia(constraints)
+                    .then((mediaStream) => {
+                        webcamVideo.srcObject = mediaStream;
+                        webcamVideo.onloadedmetadata = () => {
+                            webcamVideo.play();
+                            webcamVideo.addEventListener("timeupdate", videoTimeUpdateHandler);
+                        };
+                    })
+                    .catch((err) => {
+                        console.error(`${err.name}: ${err.message}`);
+                    });
+            }
+
+
+            async function stop() {
+                websocket.close();
+                navigator.mediaDevices.getUserMedia({ video: true }).then((mediaStream) => {
+                    mediaStream.getTracks().forEach((track) => track.stop());
+                });
+                webcamVideo.removeEventListener("timeupdate", videoTimeUpdateHandler);
+                webcamsEl.removeEventListener("change", switchCamera);
+                webcamVideo.srcObject = null;
+            }
+            return {
+                start,
+                stop
+            }
+        }
+        function toggleMessage(type) {
+            errorEl.hidden = false;
+            errorEl.scrollIntoView();
+            switch (type) {
+                case "error":
+                    errorEl.innerText = "To many users are using the same GPU, please try again later.";
+                    errorEl.classList.toggle("bg-red-300", "text-red-900");
+                    break;
+                case "success":
+                    errorEl.innerText = "Your session has ended, please start a new one.";
+                    errorEl.classList.toggle("bg-green-300", "text-green-900");
+                    break;
+            }
+            setTimeout(() => {
+                errorEl.hidden = true;
+            }, 2000);
+        }
+        function snapImage() {
+            try {
+                const zeroth = {};
+                const exif = {};
+                const gps = {};
+                zeroth[piexif.ImageIFD.Make] = "LCM Image-to-Image ControNet";
+                zeroth[piexif.ImageIFD.ImageDescription] = `prompt: ${getValue("#prompt")} | seed: ${getValue("#seed")} | guidance_scale: ${getValue("#guidance-scale")} | strength: ${getValue("#strength")} | controlnet_start: ${getValue("#controlnet_start")} | controlnet_end: ${getValue("#controlnet_end")} | lcm_steps: ${getValue("#lcm_steps")} | steps: ${getValue("#steps")}`;
+                zeroth[piexif.ImageIFD.Software] = "https://github.com/radames/Real-Time-Latent-Consistency-Model";
+                exif[piexif.ExifIFD.DateTimeOriginal] = new Date().toISOString();
+
+                const exifObj = { "0th": zeroth, "Exif": exif, "GPS": gps };
+                const exifBytes = piexif.dump(exifObj);
+
+                const canvas = document.createElement("canvas");
+                canvas.width = imageEl.naturalWidth;
+                canvas.height = imageEl.naturalHeight;
+                const ctx = canvas.getContext("2d");
+                ctx.drawImage(imageEl, 0, 0);
+                const dataURL = canvas.toDataURL("image/jpeg");
+                const withExif = piexif.insert(exifBytes, dataURL);
+
+                const a = document.createElement("a");
+                a.href = withExif;
+                a.download = `lcm_txt_2_img${Date.now()}.png`;
+                a.click();
+            } catch (err) {
+                console.log(err);
+            }
+        }
+
+
+        const lcmLive = LCMLive(videoEl, imageEl);
+        startBtn.addEventListener("click", async () => {
+            try {
+                startBtn.disabled = true;
+                snapBtn.disabled = false;
+                const res = await lcmLive.start();
+                startBtn.disabled = false;
+                if (res.status === "timeout")
+                    toggleMessage("success")
+            } catch (err) {
+                console.log(err);
+                toggleMessage("error")
+                startBtn.disabled = false;
+            }
+        });
+        stopBtn.addEventListener("click", () => {
+            lcmLive.stop();
+        });
+        window.addEventListener("beforeunload", () => {
+            lcmLive.stop();
+        });
+        snapBtn.addEventListener("click", snapImage);
+        setInterval(() =>
+            fetch("/queue_size")
+                .then((res) => res.json())
+                .then((data) => {
+                    queueSizeEl.innerText = data.queue_size;
+                })
+                .catch((err) => {
+                    console.log(err);
+                })
+            , 5000);
+    </script>
+</head>
+
+<body class="text-black dark:bg-gray-900 dark:text-white">
+    <div class="fixed right-2 top-2 p-4 font-bold text-sm rounded-lg max-w-xs text-center" id="error">
+    </div>
+    <main class="container mx-auto px-4 py-4 max-w-4xl flex flex-col gap-4">
+        <article class="text-center max-w-xl mx-auto">
+            <h1 class="text-3xl font-bold">Real-Time Latent Consistency Model</h1>
+            <h2 class="text-2xl font-bold mb-4">ControlNet</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="flex text-normal px-1 py-1 border border-gray-700 rounded-md items-center">
+                <textarea type="text" id="prompt" class="font-light w-full px-3 py-2 mx-1  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="font-medium cursor-pointer">Advanced Options</summary>
+                <div class="grid grid-cols-3 sm:grid-cols-6 items-center gap-3 py-3">
+                    <label for="webcams" class="text-sm font-medium">Camera Options: </label>
+                    <select id="webcams" class="text-sm border-2 border-gray-500 rounded-md font-light dark:text-black">
+                    </select>
+                    <div></div>
+                    <label class="text-sm font-medium " for="steps">Inference Steps
+                    </label>
+                    <input type="range" id="steps" name="steps" min="1" max="20" value="4"
+                        oninput="this.nextElementSibling.value = Number(this.value)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        4</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="lcm_steps">LCM Inference Steps
+                    </label>
+                    <input type="range" id="lcm_steps" name="lcm_steps" min="2" max="60" value="50"
+                        oninput="this.nextElementSibling.value = Number(this.value)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        50</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="guidance-scale">Guidance Scale
+                    </label>
+                    <input type="range" id="guidance-scale" name="guidance-scale" min="0" max="30" step="0.001"
+                        value="8.0" oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        8.0</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="strength">Strength</label>
+                    <input type="range" id="strength" name="strength" min="0.1" max="1" step="0.001" value="0.50"
+                        oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        0.5</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="controlnet_scale">ControlNet Condition Scale</label>
+                    <input type="range" id="controlnet_scale" name="controlnet_scale" min="0.0" max="1" step="0.001"
+                        value="0.80" oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        0.8</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="controlnet_start">ControlNet Guidance Start</label>
+                    <input type="range" id="controlnet_start" name="controlnet_start" min="0.0" max="1.0" step="0.001"
+                        value="0.0" oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        0.0</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="controlnet_end">ControlNet Guidance End</label>
+                    <input type="range" id="controlnet_end" name="controlnet_end" min="0.0" max="1.0" step="0.001"
+                        value="1.0" oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        1.0</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="canny_low_threshold">Canny Low Threshold</label>
+                    <input type="range" id="canny_low_threshold" name="canny_low_threshold" min="0.0" max="1.0"
+                        step="0.001" value="0.2"
+                        oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        0.2</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="canny_high_threshold">Canny High Threshold</label>
+                    <input type="range" id="canny_high_threshold" name="canny_high_threshold" min="0.0" max="1.0"
+                        step="0.001" value="0.8"
+                        oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        0.8</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="seed">Seed</label>
+                    <input type="number" id="seed" name="seed" value="299792458"
+                        class="font-light border border-gray-700 text-right rounded-md p-2 dark:text-black">
+                    <button
+                        onclick="document.querySelector('#seed').value = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER)"
+                        class="button">
+                        Rand
+                    </button>
+                    <!--  -->
+                    <!--  -->
+                    <label class="text-sm font-medium" for="dimension">Image Dimensions</label>
+                    <div class="col-span-2 flex gap-2">
+                        <div class="flex gap-1">
+                            <input type="radio" id="dimension512" name="dimension" value="[512,512]" checked
+                                class="cursor-pointer">
+                            <label for="dimension512" class="text-sm cursor-pointer">512x512</label>
+                        </div>
+                        <div class="flex gap-1">
+                            <input type="radio" id="dimension768" name="dimension" value="[768,768]"
+                                lass="cursor-pointer">
+                            <label for="dimension768" class="text-sm cursor-pointer">768x768</label>
+                        </div>
+                    </div>
+                    <!--  -->
+                </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 rounded-lg border border-slate-300 overflow-hidden">
+            <img id="player" class="w-full aspect-square rounded-lg"
+                src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=">
+            <div class="absolute top-0 left-0 w-1/4 aspect-square">
+                <video id="webcam" class="w-full aspect-square relative z-10 object-cover" playsinline autoplay muted
+                    loop></video>
+                <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 448" width="100"
+                    class="w-full p-4 absolute top-0 opacity-20 z-0">
+                    <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>
+</body>
+
+</html>

latent_consistency_controlnet.py

@@ -0,0 +1,1094 @@
+# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers import (
+    AutoencoderKL,
+    AutoencoderTiny,
+    ConfigMixin,
+    DiffusionPipeline,
+    SchedulerMixin,
+    UNet2DConditionModel,
+    ControlNetModel,
+    logging,
+)
+from diffusers.configuration_utils import register_to_config
+from diffusers.image_processor import VaeImageProcessor, PipelineImageInput
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import (
+    StableDiffusionSafetyChecker,
+)
+from diffusers.utils import BaseOutput
+
+from diffusers.utils.torch_utils import randn_tensor, is_compiled_module
+from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
+
+
+import PIL.Image
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
+    _optional_components = ["scheduler"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        controlnet: Union[
+            ControlNetModel,
+            List[ControlNetModel],
+            Tuple[ControlNetModel],
+            MultiControlNetModel,
+        ],
+        unet: UNet2DConditionModel,
+        scheduler: "LCMScheduler",
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        scheduler = (
+            scheduler
+            if scheduler is not None
+            else LCMScheduler_X(
+                beta_start=0.00085,
+                beta_end=0.0120,
+                beta_schedule="scaled_linear",
+                prediction_type="epsilon",
+            )
+        )
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            controlnet=controlnet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.control_image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_convert_rgb=True,
+            do_normalize=False,
+        )
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        prompt_embeds: None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+        """
+
+        if prompt is not None and isinstance(prompt, str):
+            pass
+        elif prompt is not None and isinstance(prompt, list):
+            len(prompt)
+        else:
+            prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(
+                prompt, padding="longest", return_tensors="pt"
+            ).input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[
+                -1
+            ] and not torch.equal(text_input_ids, untruncated_ids):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+                )
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if (
+                hasattr(self.text_encoder.config, "use_attention_mask")
+                and self.text_encoder.config.use_attention_mask
+            ):
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            bs_embed * num_images_per_prompt, seq_len, -1
+        )
+
+        # Don't need to get uncond prompt embedding because of LCM Guided Distillation
+        return prompt_embeds
+
+    def run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            if torch.is_tensor(image):
+                feature_extractor_input = self.image_processor.postprocess(
+                    image, output_type="pil"
+                )
+            else:
+                feature_extractor_input = self.image_processor.numpy_to_pil(image)
+            safety_checker_input = self.feature_extractor(
+                feature_extractor_input, return_tensors="pt"
+            ).to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        return image, has_nsfw_concept
+
+    def prepare_control_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance=False,
+        guess_mode=False,
+    ):
+        image = self.control_image_processor.preprocess(
+            image, height=height, width=width
+        ).to(dtype=dtype)
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        if do_classifier_free_guidance and not guess_mode:
+            image = torch.cat([image] * 2)
+
+        return image
+
+    def prepare_latents(
+        self,
+        image,
+        timestep,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        latents=None,
+        generator=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        image = image.to(device=device, dtype=dtype)
+
+        # batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            elif isinstance(generator, list):
+                if isinstance(self.vae, AutoencoderTiny):
+                    init_latents = [
+                        self.vae.encode(image[i : i + 1]).latents
+                        for i in range(batch_size)
+                    ]
+                else:
+                    init_latents = [
+                        self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i])
+                        for i in range(batch_size)
+                    ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                if isinstance(self.vae, AutoencoderTiny):
+                    init_latents = self.vae.encode(image).latents
+                else:
+                    init_latents = self.vae.encode(image).latent_dist.sample(generator)
+
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        if (
+            batch_size > init_latents.shape[0]
+            and batch_size % init_latents.shape[0] == 0
+        ):
+            # expand init_latents for batch_size
+            deprecation_message = (
+                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"
+                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+                " your script to pass as many initial images as text prompts to suppress this warning."
+            )
+            # deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat(
+                [init_latents] * additional_image_per_prompt, dim=0
+            )
+        elif (
+            batch_size > init_latents.shape[0]
+            and batch_size % init_latents.shape[0] != 0
+        ):
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        shape = init_latents.shape
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        # get latents
+        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+        latents = init_latents
+
+        return latents
+
+        if latents is None:
+            latents = torch.randn(shape, dtype=dtype).to(device)
+        else:
+            latents = latents.to(device)
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32):
+        """
+        see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+        Args:
+        timesteps: torch.Tensor: generate embedding vectors at these timesteps
+        embedding_dim: int: dimension of the embeddings to generate
+        dtype: data type of the generated embeddings
+        Returns:
+        embedding vectors with shape `(len(timesteps), embedding_dim)`
+        """
+        assert len(w.shape) == 1
+        w = w * 1000.0
+
+        half_dim = embedding_dim // 2
+        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+        emb = w.to(dtype)[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+        if embedding_dim % 2 == 1:  # zero pad
+            emb = torch.nn.functional.pad(emb, (0, 1))
+        assert emb.shape == (w.shape[0], embedding_dim)
+        return emb
+
+    def get_timesteps(self, num_inference_steps, strength, device):
+        # get the original timestep using init_timestep
+        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+        t_start = max(num_inference_steps - init_timestep, 0)
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+
+        return timesteps, num_inference_steps - t_start
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        image: PipelineImageInput = None,
+        control_image: PipelineImageInput = None,
+        strength: float = 0.8,
+        height: Optional[int] = 768,
+        width: Optional[int] = 768,
+        guidance_scale: float = 7.5,
+        num_images_per_prompt: Optional[int] = 1,
+        latents: Optional[torch.FloatTensor] = None,
+        generator: Optional[torch.Generator] = None,
+        num_inference_steps: int = 4,
+        lcm_origin_steps: int = 50,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
+        guess_mode: bool = True,
+        control_guidance_start: Union[float, List[float]] = 0.0,
+        control_guidance_end: Union[float, List[float]] = 1.0,
+    ):
+        controlnet = (
+            self.controlnet._orig_mod
+            if is_compiled_module(self.controlnet)
+            else self.controlnet
+        )
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        if not isinstance(control_guidance_start, list) and isinstance(
+            control_guidance_end, list
+        ):
+            control_guidance_start = len(control_guidance_end) * [
+                control_guidance_start
+            ]
+        elif not isinstance(control_guidance_end, list) and isinstance(
+            control_guidance_start, list
+        ):
+            control_guidance_end = len(control_guidance_start) * [control_guidance_end]
+        elif not isinstance(control_guidance_start, list) and not isinstance(
+            control_guidance_end, list
+        ):
+            mult = (
+                len(controlnet.nets)
+                if isinstance(controlnet, MultiControlNetModel)
+                else 1
+            )
+            control_guidance_start, control_guidance_end = mult * [
+                control_guidance_start
+            ], mult * [control_guidance_end]
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # do_classifier_free_guidance = guidance_scale > 0.0  # In LCM Implementation:  cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG)
+        global_pool_conditions = (
+            controlnet.config.global_pool_conditions
+            if isinstance(controlnet, ControlNetModel)
+            else controlnet.nets[0].config.global_pool_conditions
+        )
+        guess_mode = guess_mode or global_pool_conditions
+        # 3. Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+        )
+
+        # 3.5 encode image
+        image = self.image_processor.preprocess(image)
+
+        if isinstance(controlnet, ControlNetModel):
+            control_image = self.prepare_control_image(
+                image=control_image,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=controlnet.dtype,
+                guess_mode=guess_mode,
+            )
+        elif isinstance(controlnet, MultiControlNetModel):
+            control_images = []
+
+            for control_image_ in control_image:
+                control_image_ = self.prepare_control_image(
+                    image=control_image_,
+                    width=width,
+                    height=height,
+                    batch_size=batch_size * num_images_per_prompt,
+                    num_images_per_prompt=num_images_per_prompt,
+                    device=device,
+                    dtype=controlnet.dtype,
+                    do_classifier_free_guidance=do_classifier_free_guidance,
+                    guess_mode=guess_mode,
+                )
+
+                control_images.append(control_image_)
+
+            control_image = control_images
+        else:
+            assert False
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(strength, num_inference_steps, lcm_origin_steps)
+        # timesteps = self.scheduler.timesteps
+        # timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, 1.0, device)
+        timesteps = self.scheduler.timesteps
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        print("timesteps: ", timesteps)
+
+        # 5. Prepare latent variable
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            image,
+            latent_timestep,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            latents,
+        )
+        bs = batch_size * num_images_per_prompt
+
+        # 6. Get Guidance Scale Embedding
+        w = torch.tensor(guidance_scale).repeat(bs)
+        w_embedding = self.get_w_embedding(w, embedding_dim=256).to(
+            device=device, dtype=latents.dtype
+        )
+        controlnet_keep = []
+        for i in range(len(timesteps)):
+            keeps = [
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
+                for s, e in zip(control_guidance_start, control_guidance_end)
+            ]
+            controlnet_keep.append(
+                keeps[0] if isinstance(controlnet, ControlNetModel) else keeps
+            )
+        # 7. LCM MultiStep Sampling Loop:
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                ts = torch.full((bs,), t, device=device, dtype=torch.long)
+                latents = latents.to(prompt_embeds.dtype)
+                if guess_mode:
+                    # Infer ControlNet only for the conditional batch.
+                    control_model_input = latents
+                    control_model_input = self.scheduler.scale_model_input(
+                        control_model_input, ts
+                    )
+                    controlnet_prompt_embeds = prompt_embeds
+                else:
+                    control_model_input = latents
+                    controlnet_prompt_embeds = prompt_embeds
+                if isinstance(controlnet_keep[i], list):
+                    cond_scale = [
+                        c * s
+                        for c, s in zip(
+                            controlnet_conditioning_scale, controlnet_keep[i]
+                        )
+                    ]
+                else:
+                    controlnet_cond_scale = controlnet_conditioning_scale
+                    if isinstance(controlnet_cond_scale, list):
+                        controlnet_cond_scale = controlnet_cond_scale[0]
+                    cond_scale = controlnet_cond_scale * controlnet_keep[i]
+
+                down_block_res_samples, mid_block_res_sample = self.controlnet(
+                    control_model_input,
+                    ts,
+                    encoder_hidden_states=controlnet_prompt_embeds,
+                    controlnet_cond=control_image,
+                    conditioning_scale=cond_scale,
+                    guess_mode=guess_mode,
+                    return_dict=False,
+                )
+                # model prediction (v-prediction, eps, x)
+                model_pred = self.unet(
+                    latents,
+                    ts,
+                    timestep_cond=w_embedding,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    down_block_additional_residuals=down_block_res_samples,
+                    mid_block_additional_residual=mid_block_res_sample,
+                    return_dict=False,
+                )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents, denoised = self.scheduler.step(
+                    model_pred, i, t, latents, return_dict=False
+                )
+
+                # # call the callback, if provided
+                # if i == len(timesteps) - 1:
+                progress_bar.update()
+
+        denoised = denoised.to(prompt_embeds.dtype)
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.unet.to("cpu")
+            self.controlnet.to("cpu")
+            torch.cuda.empty_cache()
+        if not output_type == "latent":
+            image = self.vae.decode(
+                denoised / self.vae.config.scaling_factor, return_dict=False
+            )[0]
+            image, has_nsfw_concept = self.run_safety_checker(
+                image, device, prompt_embeds.dtype
+            )
+        else:
+            image = denoised
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(
+            image, output_type=output_type, do_denormalize=do_denormalize
+        )
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(
+            images=image, nsfw_content_detected=has_nsfw_concept
+        )
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class LCMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    denoised: Optional[torch.FloatTensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+    Args:
+        betas (`torch.FloatTensor`):
+            the betas that the scheduler is being initialized with.
+    Returns:
+        `torch.FloatTensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class LCMScheduler_X(SchedulerMixin, ConfigMixin):
+    """
+    `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps. You can use a combination of `offset=1` and
+            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
+            Diffusion.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    # _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        clip_sample: bool = True,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        clip_sample_range: float = 1.0,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading",
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(
+                beta_start, beta_end, num_train_timesteps, dtype=torch.float32
+            )
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = (
+                torch.linspace(
+                    beta_start**0.5,
+                    beta_end**0.5,
+                    num_train_timesteps,
+                    dtype=torch.float32,
+                )
+                ** 2
+            )
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(
+                f"{beta_schedule} does is not implemented for {self.__class__}"
+            )
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = (
+            torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+        )
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(
+            np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)
+        )
+
+    def scale_model_input(
+        self, sample: torch.FloatTensor, timestep: Optional[int] = None
+    ) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def _get_variance(self, timestep, prev_timestep):
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = (
+            self.alphas_cumprod[prev_timestep]
+            if prev_timestep >= 0
+            else self.final_alpha_cumprod
+        )
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        variance = (beta_prod_t_prev / beta_prod_t) * (
+            1 - alpha_prod_t / alpha_prod_t_prev
+        )
+
+        return variance
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, height, width = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = (
+                sample.float()
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * height * width)
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = (
+            torch.clamp(sample, -s, s) / s
+        )  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, height, width)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(
+        self,
+        stength,
+        num_inference_steps: int,
+        lcm_origin_steps: int,
+        device: Union[str, torch.device] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+
+        # LCM Timesteps Setting:  # Linear Spacing
+        c = self.config.num_train_timesteps // lcm_origin_steps
+        lcm_origin_timesteps = (
+            np.asarray(list(range(1, int(lcm_origin_steps * stength) + 1))) * c - 1
+        )  # LCM Training  Steps Schedule
+        skipping_step = max(len(lcm_origin_timesteps) // num_inference_steps, 1)
+        timesteps = lcm_origin_timesteps[::-skipping_step][
+            :num_inference_steps
+        ]  # LCM Inference Steps Schedule
+
+        self.timesteps = torch.from_numpy(timesteps.copy()).to(device)
+
+    def get_scalings_for_boundary_condition_discrete(self, t):
+        self.sigma_data = 0.5  # Default: 0.5
+
+        # By dividing 0.1: This is almost a delta function at t=0.
+        c_skip = self.sigma_data**2 / ((t / 0.1) ** 2 + self.sigma_data**2)
+        c_out = (t / 0.1) / ((t / 0.1) ** 2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timeindex: int,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.FloatTensor] = None,
+        return_dict: bool = True,
+    ) -> Union[LCMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            eta (`float`):
+                The weight of noise for added noise in diffusion step.
+            use_clipped_model_output (`bool`, defaults to `False`):
+                If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
+                because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
+                clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
+                `use_clipped_model_output` has no effect.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.FloatTensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`CycleDiffusion`].
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        # 1. get previous step value
+        prev_timeindex = timeindex + 1
+        if prev_timeindex < len(self.timesteps):
+            prev_timestep = self.timesteps[prev_timeindex]
+        else:
+            prev_timestep = timestep
+
+        # 2. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = (
+            self.alphas_cumprod[prev_timestep]
+            if prev_timestep >= 0
+            else self.final_alpha_cumprod
+        )
+
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        # 3. Get scalings for boundary conditions
+        c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
+
+        # 4. Different Parameterization:
+        parameterization = self.config.prediction_type
+
+        if parameterization == "epsilon":  # noise-prediction
+            pred_x0 = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
+
+        elif parameterization == "sample":  # x-prediction
+            pred_x0 = model_output
+
+        elif parameterization == "v_prediction":  # v-prediction
+            pred_x0 = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
+
+        # 4. Denoise model output using boundary conditions
+        denoised = c_out * pred_x0 + c_skip * sample
+
+        # 5. Sample z ~ N(0, I), For MultiStep Inference
+        # Noise is not used for one-step sampling.
+        if len(self.timesteps) > 1:
+            noise = torch.randn(model_output.shape).to(model_output.device)
+            prev_sample = (
+                alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
+            )
+        else:
+            prev_sample = denoised
+
+        if not return_dict:
+            return (prev_sample, denoised)
+
+        return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        alphas_cumprod = self.alphas_cumprod.to(
+            device=original_samples.device, dtype=original_samples.dtype
+        )
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = (
+            sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        )
+        return noisy_samples
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(
+        self,
+        sample: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        alphas_cumprod = self.alphas_cumprod.to(
+            device=sample.device, dtype=sample.dtype
+        )
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

requirements.txt

@@ -7,4 +7,5 @@ fastapi==0.104.0
 uvicorn==0.23.2
 Pillow==10.1.0
 accelerate==0.24.0
-compel==2.0.2
+compel==2.0.2
+controlnet-aux==0.0.7