refactor: Clean code and refactor app to use torch.hub

app.py

@@ -1,46 +1,19 @@
-# Copyright 2024 Anton Obukhov, ETH Zurich. 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.
-# --------------------------------------------------------------------------
-# If you find this code useful, we kindly ask you to cite our paper in your work.
-# Please find bibtex at: https://github.com/prs-eth/Marigold#-citation
-# More information about the method can be found at https://marigoldmonodepth.github.io
-# --------------------------------------------------------------------------
 from __future__ import annotations
 
 import functools
 import os
 import tempfile
-
-import diffusers
+import torch
 import gradio as gr
-import imageio as imageio
-import numpy as np
-import spaces
-import torch as torch
-torch.backends.cuda.matmul.allow_tf32 = True
 from PIL import Image
 from gradio_imageslider import ImageSlider
-from tqdm import tqdm
-
 from pathlib import Path
-import gradio
 from gradio.utils import get_cache_folder
-from stablenormal.pipeline_yoso_normal import YOSONormalsPipeline
-from stablenormal.pipeline_stablenormal import StableNormalPipeline
-from stablenormal.scheduler.heuristics_ddimsampler import HEURI_DDIMScheduler
 
-class Examples(gradio.helpers.Examples):
+# Constants
+DEFAULT_SHARPNESS = 2
+
+class Examples(gr.helpers.Examples):
     def __init__(self, *args, directory_name=None, **kwargs):
         super().__init__(*args, **kwargs, _initiated_directly=False)
         if directory_name is not None:
@@ -48,250 +21,94 @@ class Examples(gradio.helpers.Examples):
             self.cached_file = Path(self.cached_folder) / "log.csv"
         self.create()
 
-
-default_seed = 2024
-default_batch_size = 1
-
-default_image_processing_resolution = 768
-
-default_video_num_inference_steps = 10
-default_video_processing_resolution = 768
-default_video_out_max_frames = 60
-
-def process_image_check(path_input):
-    if path_input is None:
-        raise gr.Error(
-            "Missing image in the first pane: upload a file or use one from the gallery below."
-        )
-
-def resize_image(input_image, resolution):
-    # Ensure input_image is a PIL Image object
-    if not isinstance(input_image, Image.Image):
-        raise ValueError("input_image should be a PIL Image object")
-
-    # Convert image to numpy array
-    input_image_np = np.asarray(input_image)
-
-    # Get image dimensions
-    H, W, C = input_image_np.shape
-    H = float(H)
-    W = float(W)
-    
-    # Calculate the scaling factor
-    k = float(resolution) / min(H, W)
-    
-    # Determine new dimensions
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    
-    # Resize the image using PIL's resize method
-    img = input_image.resize((W, H), Image.Resampling.LANCZOS)
-    
-    return img
+def load_predictor():
+    """Load model predictor using torch.hub"""
+    predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal", trust_repo=True, 
+                           local_cache_dir='./weights')
+    return predictor
 
 def process_image(
-    pipe,
-    path_input,
-):
-    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
-    print(f"Processing image {name_base}{name_ext}")
+    predictor,
+    path_input: str,
+    sharpness: int = DEFAULT_SHARPNESS,
+    data_type: str = "object"
+) -> tuple:
+    """Process single image"""
+    if path_input is None:
+        raise gr.Error("Please upload an image or select one from the gallery.")
+        
+    name_base = os.path.splitext(os.path.basename(path_input))[0]
+    out_path = os.path.join(tempfile.mkdtemp(), f"{name_base}_normal.png")
 
-    path_output_dir = tempfile.mkdtemp()
-    path_out_png = os.path.join(path_output_dir, f"{name_base}_normal_colored.png")
+    # Load and process image
     input_image = Image.open(path_input)
-    input_image = resize_image(input_image, default_image_processing_resolution)
+    normal_image = predictor(input_image, num_inference_steps=sharpness, 
+                             match_input_resolution=False, data_type=data_type)
+    normal_image.save(out_path)
 
-    pipe_out = pipe(
-        input_image,
-        match_input_resolution=False,
-        processing_resolution=max(input_image.size)
-    )
-
-    normal_pred = pipe_out.prediction[0, :, :]
-    normal_colored = pipe.image_processor.visualize_normals(pipe_out.prediction)
-    normal_colored[-1].save(path_out_png)
-    yield [input_image, path_out_png]
+    yield [input_image, out_path]
 
-def center_crop(img):
-    # Open the image file
-    img_width, img_height = img.size
-    crop_width =min(img_width, img_height)
-    # Calculate the cropping box
-    left = (img_width - crop_width) / 2
-    top = (img_height - crop_width) / 2
-    right = (img_width + crop_width) / 2
-    bottom = (img_height + crop_width) / 2
+def create_demo():
+    # Load model
+    predictor = load_predictor()
     
-    # Crop the image
-    img_cropped = img.crop((left, top, right, bottom))
-    return img_cropped
-
-def process_video(
-    pipe,
-    path_input,
-    out_max_frames=default_video_out_max_frames,
-    target_fps=10,
-    progress=gr.Progress(),
-):
-    if path_input is None:
-        raise gr.Error(
-            "Missing video in the first pane: upload a file or use one from the gallery below."
-        )
-
-    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
-    print(f"Processing video {name_base}{name_ext}")
-
-    path_output_dir = tempfile.mkdtemp()
-    path_out_vis = os.path.join(path_output_dir, f"{name_base}_normal_colored.mp4")
-
-    init_latents = None
-    reader, writer = None, None
-    try:
-        reader = imageio.get_reader(path_input)
-
-        meta_data = reader.get_meta_data()
-        fps = meta_data["fps"]
-        size = meta_data["size"]
-        duration_sec = meta_data["duration"]
-
-        writer = imageio.get_writer(path_out_vis, fps=target_fps)
-
-        out_frame_id = 0
-        pbar = tqdm(desc="Processing Video", total=duration_sec)
-
-        for frame_id, frame in enumerate(reader):
-            if frame_id % (fps // target_fps) != 0:
-                continue
-            else:
-                out_frame_id += 1
-                pbar.update(1)
-            if out_frame_id > out_max_frames:
-                break
-
-            frame_pil = Image.fromarray(frame)
-            frame_pil = center_crop(frame_pil)
-            pipe_out = pipe(
-                frame_pil,
-                match_input_resolution=False,
-                latents=init_latents
-            )
-
-            if init_latents is None:
-                init_latents = pipe_out.gaus_noise
-            processed_frame = pipe.image_processor.visualize_normals(  # noqa
-                pipe_out.prediction
-            )[0]
-            processed_frame = np.array(processed_frame)
-
-            _processed_frame = imageio.core.util.Array(processed_frame)
-            writer.append_data(_processed_frame)
-            
-            yield (
-                [frame_pil, processed_frame],
-                None,
-            )
-    finally:
-
-        if writer is not None:
-            writer.close()
-
-        if reader is not None:
-            reader.close()
-
-    yield (
-        [frame_pil, processed_frame],
-        [path_out_vis,]
-    )
-
-
-def run_demo_server(pipe):
-    process_pipe_image = spaces.GPU(functools.partial(process_image, pipe))
-    process_pipe_video = spaces.GPU(
-        functools.partial(process_video, pipe), duration=120
-    )
-
-    gradio_theme = gr.themes.Default()
-
-    with gr.Blocks(
-        theme=gradio_theme,
+    # Create processing functions for each data type
+    process_object = functools.partial(process_image, predictor, data_type="object")
+    process_scene = functools.partial(process_image, predictor, data_type="indoor")
+    process_human = functools.partial(process_image, predictor, data_type="object")
+
+    # Define markdown content
+    HEADER_MD = """
+    # StableNormal: Reducing Diffusion Variance for Stable and Sharp Normal
+    <p align="center">
+    <a title="Website" href="https://stable-x.github.io/StableNormal/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+        <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+    </a>
+    <a title="arXiv" href="https://arxiv.org/abs/2406.16864" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+    </a>
+    <a title="Github" href="https://github.com/Stable-X/StableNormal" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+        <img src="https://img.shields.io/github/stars/Stable-X/StableDelight?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+    </a>
+    <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+        <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+    </a>
+    """
+
+    # Create interface
+    demo = gr.Blocks(
         title="Stable Normal Estimation",
         css="""
-            #download {
-                height: 118px;
-            }
-            .slider .inner {
-                width: 5px;
-                background: #FFF;
-            }
-            .viewport {
-                aspect-ratio: 4/3;
-            }
-            .tabs button.selected {
-                font-size: 20px !important;
-                color: crimson !important;
-            }
-            h1 {
-                text-align: center;
-                display: block;
-            }
-            h2 {
-                text-align: center;
-                display: block;
-            }
-            h3 {
-                text-align: center;
-                display: block;
-            }
-            .md_feedback li {
-                margin-bottom: 0px !important;
-            }
-        """,
-        head="""
-            <script async src="https://www.googletagmanager.com/gtag/js?id=G-1FWSVCGZTG"></script>
-            <script>
-                window.dataLayer = window.dataLayer || [];
-                function gtag() {dataLayer.push(arguments);}
-                gtag('js', new Date());
-                gtag('config', 'G-1FWSVCGZTG');
-            </script>
-        """,
-    ) as demo:
-        gr.Markdown(
-            """
-            # StableNormal: Reducing Diffusion Variance for Stable and Sharp Normal
-            <p align="center">
-
-            <a title="Website" href="https://stable-x.github.io/StableNormal/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
-            </a>
-            <a title="arXiv" href="https://arxiv.org/abs/2406.16864" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
-            </a>
-            <a title="Github" href="https://github.com/Stable-X/StableNormal" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://img.shields.io/github/stars/Stable-X/StableDelight?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
-            </a>
-            <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
-            </a>
+            .slider .inner { width: 5px; background: #FFF; }
+            .viewport { aspect-ratio: 4/3; }
+            .tabs button.selected { font-size: 20px !important; color: crimson !important; }
+            h1, h2, h3 { text-align: center; display: block; }
+            .md_feedback li { margin-bottom: 0px !important; }
         """
-        )
-        with gr.Tabs(elem_classes=["tabs"]):
-            with gr.Tab("Image"):
+    )
+
+    with demo:
+        gr.Markdown(HEADER_MD)
+
+        with gr.Tabs() as tabs:
+            # Object Tab
+            with gr.Tab("Object"):
                 with gr.Row():
                     with gr.Column():
-                        image_input = gr.Image(
-                            label="Input Image",
-                            type="filepath",
+                        object_input = gr.Image(label="Input Object Image", type="filepath")
+                        object_sharpness = gr.Slider(
+                            minimum=1,
+                            maximum=10,
+                            value=DEFAULT_SHARPNESS,
+                            step=1,
+                            label="Sharpness (inference steps)",
+                            info="Higher values produce sharper results but take longer"
                         )
                         with gr.Row():
-                            image_submit_btn = gr.Button(
-                                value="Compute Normal", variant="primary"
-                            )
-                            image_reset_btn = gr.Button(value="Reset")
+                            object_submit_btn = gr.Button("Compute Normal", variant="primary")
+                            object_reset_btn = gr.Button("Reset")
                     with gr.Column():
-                        image_output_slider = ImageSlider(
+                        object_output_slider = ImageSlider(
                             label="Normal outputs",
                             type="filepath",
                             show_download_button=True,
@@ -302,32 +119,38 @@ def run_demo_server(pipe):
                         )
 
                 Examples(
-                    fn=process_pipe_image,
+                    fn=process_object,
                     examples=sorted([
-                        os.path.join("files", "image", name)
-                        for name in os.listdir(os.path.join("files", "image"))
+                        os.path.join("files", "object", name)
+                        for name in os.listdir(os.path.join("files", "object"))
+                        if os.path.exists(os.path.join("files", "object"))
                     ]),
-                    inputs=[image_input],
-                    outputs=[image_output_slider],
+                    inputs=[object_input],
+                    outputs=[object_output_slider],
                     cache_examples=True,
-                    directory_name="examples_image",
+                    directory_name="examples_object",
+                    examples_per_page=50,
                 )
 
-            with gr.Tab("Video"):
+            # Scene Tab
+            with gr.Tab("Scene"):
                 with gr.Row():
                     with gr.Column():
-                        video_input = gr.Video(
-                            label="Input Video",
-                            sources=["upload", "webcam"],
+                        scene_input = gr.Image(label="Input Scene Image", type="filepath")
+                        scene_sharpness = gr.Slider(
+                            minimum=1,
+                            maximum=10,
+                            value=DEFAULT_SHARPNESS,
+                            step=1,
+                            label="Sharpness (inference steps)",
+                            info="Higher values produce sharper results but take longer"
                         )
                         with gr.Row():
-                            video_submit_btn = gr.Button(
-                                value="Compute Normal", variant="primary"
-                            )
-                            video_reset_btn = gr.Button(value="Reset")
+                            scene_submit_btn = gr.Button("Compute Normal", variant="primary")
+                            scene_reset_btn = gr.Button("Reset")
                     with gr.Column():
-                        processed_frames = ImageSlider(
-                            label="Realtime Visualization",
+                        scene_output_slider = ImageSlider(
+                            label="Normal outputs",
                             type="filepath",
                             show_download_button=True,
                             show_share_button=True,
@@ -335,111 +158,127 @@ def run_demo_server(pipe):
                             elem_classes="slider",
                             position=0.25,
                         )
-                        video_output_files = gr.Files(
+
+                Examples(
+                    fn=process_scene,
+                    examples=sorted([
+                        os.path.join("files", "scene", name)
+                        for name in os.listdir(os.path.join("files", "scene"))
+                        if os.path.exists(os.path.join("files", "scene"))
+                    ]),
+                    inputs=[scene_input],
+                    outputs=[scene_output_slider],
+                    cache_examples=True,
+                    directory_name="examples_scene",
+                    examples_per_page=50,
+                )
+
+            # Human Tab
+            with gr.Tab("Human"):
+                with gr.Row():
+                    with gr.Column():
+                        human_input = gr.Image(label="Input Human Image", type="filepath")
+                        human_sharpness = gr.Slider(
+                            minimum=1,
+                            maximum=10,
+                            value=DEFAULT_SHARPNESS,
+                            step=1,
+                            label="Sharpness (inference steps)",
+                            info="Higher values produce sharper results but take longer"
+                        )
+                        with gr.Row():
+                            human_submit_btn = gr.Button("Compute Normal", variant="primary")
+                            human_reset_btn = gr.Button("Reset")
+                    with gr.Column():
+                        human_output_slider = ImageSlider(
                             label="Normal outputs",
-                            elem_id="download",
+                            type="filepath",
+                            show_download_button=True,
+                            show_share_button=True,
                             interactive=False,
+                            elem_classes="slider",
+                            position=0.25,
                         )
+
                 Examples(
-                    fn=process_pipe_video,
+                    fn=process_human,
                     examples=sorted([
-                        os.path.join("files", "video", name)
-                        for name in os.listdir(os.path.join("files", "video"))
+                        os.path.join("files", "human", name)
+                        for name in os.listdir(os.path.join("files", "human"))
+                        if os.path.exists(os.path.join("files", "human"))
                     ]),
-                    inputs=[video_input],
-                    outputs=[processed_frames, video_output_files],
-                    directory_name="examples_video",
-                    cache_examples=False,
+                    inputs=[human_input],
+                    outputs=[human_output_slider],
+                    cache_examples=True,
+                    directory_name="examples_human",
+                    examples_per_page=50,
                 )
-                
-            with gr.Tab("Panorama"):
-                with gr.Column():
-                    gr.Markdown("Coming soon")
 
-            with gr.Tab("4K Image"):
-                with gr.Column():
-                    gr.Markdown("Coming soon")
-
-        ### Image tab
-        image_submit_btn.click(
-            fn=process_image_check,
-            inputs=image_input,
+        # Event Handlers for Object Tab
+        object_submit_btn.click(
+            fn=lambda x, _: None if x else gr.Error("Please upload an image"),
+            inputs=[object_input, object_sharpness],
             outputs=None,
-            preprocess=False,
             queue=False,
         ).success(
-            fn=process_pipe_image,
-            inputs=[
-                image_input,
-            ],
-            outputs=[image_output_slider],
-            concurrency_limit=1,
+            fn=process_object,
+            inputs=[object_input, object_sharpness],
+            outputs=[object_output_slider],
         )
 
-        image_reset_btn.click(
-            fn=lambda: (
-                None,
-                None,
-                None,
-            ),
+        object_reset_btn.click(
+            fn=lambda: (None, DEFAULT_SHARPNESS, None),
             inputs=[],
-            outputs=[
-                image_input,
-                image_output_slider,
-            ],
+            outputs=[object_input, object_sharpness, object_output_slider],
             queue=False,
         )
 
-        ### Video tab
-
-        video_submit_btn.click(
-            fn=process_pipe_video,
-            inputs=[video_input],
-            outputs=[processed_frames, video_output_files],
-            concurrency_limit=1,
+        # Event Handlers for Scene Tab
+        scene_submit_btn.click(
+            fn=lambda x, _: None if x else gr.Error("Please upload an image"),
+            inputs=[scene_input, scene_sharpness],
+            outputs=None,
+            queue=False,
+        ).success(
+            fn=process_scene,
+            inputs=[scene_input, scene_sharpness],
+            outputs=[scene_output_slider],
         )
 
-        video_reset_btn.click(
-            fn=lambda: (None, None, None),
+        scene_reset_btn.click(
+            fn=lambda: (None, DEFAULT_SHARPNESS, None),
             inputs=[],
-            outputs=[video_input, processed_frames, video_output_files],
-            concurrency_limit=1,
+            outputs=[scene_input, scene_sharpness, scene_output_slider],
+            queue=False,
         )
 
-        ### Server launch
+        # Event Handlers for Human Tab
+        human_submit_btn.click(
+            fn=lambda x, _: None if x else gr.Error("Please upload an image"),
+            inputs=[human_input, human_sharpness],
+            outputs=None,
+            queue=False,
+        ).success(
+            fn=process_human,
+            inputs=[human_input, human_sharpness],
+            outputs=[human_output_slider],
+        )
 
-        demo.queue(
-            api_open=False,
-        ).launch(
-            server_name="0.0.0.0",
-            server_port=7860,
+        human_reset_btn.click(
+            fn=lambda: (None, DEFAULT_SHARPNESS, None),
+            inputs=[],
+            outputs=[human_input, human_sharpness, human_output_slider],
+            queue=False,
         )
 
+    return demo
 
 def main():
-    os.system("pip freeze")
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    x_start_pipeline = YOSONormalsPipeline.from_pretrained(
-        'Stable-X/yoso-normal-v0-3', trust_remote_code=True, variant="fp16", torch_dtype=torch.float16).to(device)
-    pipe = StableNormalPipeline.from_pretrained('Stable-X/stable-normal-v0-1', trust_remote_code=True,
-                                                variant="fp16", torch_dtype=torch.float16,
-                                                scheduler=HEURI_DDIMScheduler(prediction_type='sample', 
-                                                                              beta_start=0.00085, beta_end=0.0120, 
-                                                                              beta_schedule = "scaled_linear"))
-    pipe.x_start_pipeline = x_start_pipeline
-    pipe.to(device)
-    pipe.prior.to(device, torch.float16)
-    
-    try:
-        import xformers
-        pipe.enable_xformers_memory_efficient_attention()
-    except:
-        pass  # run without xformers
-
-    run_demo_server(pipe)
-
+    demo = create_demo()
+    demo.queue(api_open=False).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+    )
 
 if __name__ == "__main__":
-    main()
+    main()

setup.py

@@ -1,9 +0,0 @@
-from pathlib import Path
-from setuptools import setup, find_packages
-
-setup_path = Path(__file__).parent
-
-setup(
-    name = "stablenormal",
-    packages=find_packages()
-)

stablenormal/pipeline_stablenormal.py

@@ -1,1279 +0,0 @@
-# Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved.
-# Copyright 2024 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.
-# --------------------------------------------------------------------------
-# More information and citation instructions are available on the
-# --------------------------------------------------------------------------
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from PIL import Image
-from tqdm.auto import tqdm
-from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
-
-
-from diffusers.image_processor import PipelineImageInput
-from diffusers.models import (
-    AutoencoderKL,
-    UNet2DConditionModel,
-    ControlNetModel,
-)
-from diffusers.schedulers import (
-    DDIMScheduler
-)
-
-from diffusers.utils import (
-    BaseOutput,
-    logging,
-    replace_example_docstring,
-)
-
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
-
-from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
-
-
-
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.controlnet import StableDiffusionControlNetPipeline
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.marigold.marigold_image_processing import MarigoldImageProcessor
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-import torch.nn.functional as F
-
-import pdb
-
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
-Examples:
-```py
->>> import diffusers
->>> import torch
-
->>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
-...     "prs-eth/marigold-normals-lcm-v0-1", variant="fp16", torch_dtype=torch.float16
-... ).to("cuda")
-
->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
->>> normals = pipe(image)
-
->>> vis = pipe.image_processor.visualize_normals(normals.prediction)
->>> vis[0].save("einstein_normals.png")
-```
-"""
-
-
-@dataclass
-class StableNormalOutput(BaseOutput):
-    """
-    Output class for Marigold monocular normals prediction pipeline.
-
-    Args:
-        prediction (`np.ndarray`, `torch.Tensor`):
-            Predicted normals with values in the range [-1, 1]. The shape is always $numimages \times 3 \times height
-            \times width$, regardless of whether the images were passed as a 4D array or a list.
-        uncertainty (`None`, `np.ndarray`, `torch.Tensor`):
-            Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages
-            \times 1 \times height \times width$.
-        latent (`None`, `torch.Tensor`):
-            Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
-            The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$.
-    """
-
-    prediction: Union[np.ndarray, torch.Tensor]
-    latent: Union[None, torch.Tensor]
-    gaus_noise: Union[None, torch.Tensor]
-
-from einops import rearrange  
-class DINOv2_Encoder(torch.nn.Module):
-    IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406]
-    IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225]
-
-    def __init__(
-        self,
-        model_name = 'dinov2_vitl14',
-        freeze = True,
-        antialias=True,
-        device="cuda",
-        size = 448,
-    ):
-        super(DINOv2_Encoder, self).__init__()
-        
-        self.model = torch.hub.load('facebookresearch/dinov2', model_name)
-        self.model.eval().to(device)
-        self.device = device
-        self.antialias = antialias
-        self.dtype = torch.float32
-
-        self.mean = torch.Tensor(self.IMAGENET_DEFAULT_MEAN)
-        self.std = torch.Tensor(self.IMAGENET_DEFAULT_STD)
-        self.size = size
-        if freeze:
-            self.freeze()
-
-
-    def freeze(self):
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-    @torch.no_grad()
-    def encoder(self, x):
-        '''
-        x: [b h w c], range from (-1, 1), rbg
-        '''
-
-        x = self.preprocess(x).to(self.device, self.dtype)
-
-        b, c, h, w = x.shape
-        patch_h, patch_w = h // 14, w // 14
-
-        embeddings = self.model.forward_features(x)['x_norm_patchtokens']
-        embeddings = rearrange(embeddings, 'b (h w) c -> b h w c', h = patch_h, w = patch_w)
-
-        return  rearrange(embeddings, 'b h w c -> b c h w')
-
-    def preprocess(self, x):
-        ''' x
-        '''
-        # normalize to [0,1],
-        x = torch.nn.functional.interpolate(
-            x,
-            size=(self.size, self.size),
-            mode='bicubic',
-            align_corners=True,
-            antialias=self.antialias,
-        )
-
-        x = (x + 1.0) / 2.0
-        # renormalize according to dino
-        mean = self.mean.view(1, 3, 1, 1).to(x.device)
-        std = self.std.view(1, 3, 1, 1).to(x.device)
-        x = (x - mean) / std
-
-        return x
-    
-    def to(self, device, dtype=None):
-        if dtype is not None:
-            self.dtype = dtype
-            self.model.to(device, dtype)
-            self.mean.to(device, dtype)
-            self.std.to(device, dtype)
-        else:
-            self.model.to(device)
-            self.mean.to(device)
-            self.std.to(device)
-        return self
-
-    def __call__(self, x, **kwargs):
-        return self.encoder(x, **kwargs)
-
-class StableNormalPipeline(StableDiffusionControlNetPipeline):
-    """ Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io.
-    Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    The pipeline also inherits the following loading methods:
-        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
-        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
-        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
-        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
-        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
-
-    Args:
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
-        text_encoder ([`~transformers.CLIPTextModel`]):
-            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
-        tokenizer ([`~transformers.CLIPTokenizer`]):
-            A `CLIPTokenizer` to tokenize text.
-        unet ([`UNet2DConditionModel`]):
-            A `UNet2DConditionModel` to denoise the encoded image latents.
-        controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
-            Provides additional conditioning to the `unet` during the denoising process. If you set multiple
-            ControlNets as a list, the outputs from each ControlNet are added together to create one combined
-            additional conditioning.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
-            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
-        safety_checker ([`StableDiffusionSafetyChecker`]):
-            Classification module that estimates whether generated images could be considered offensive or harmful.
-            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
-            about a model's potential harms.
-        feature_extractor ([`~transformers.CLIPImageProcessor`]):
-            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
-    """
-
-    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
-    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
-    _exclude_from_cpu_offload = ["safety_checker"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        tokenizer: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel]],
-        dino_controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel]],
-        scheduler: Union[DDIMScheduler],
-        safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
-        image_encoder: CLIPVisionModelWithProjection = None,
-        requires_safety_checker: bool = True,
-        default_denoising_steps: Optional[int] = 10,
-        default_processing_resolution: Optional[int] = 768,
-        prompt="The normal map",
-        empty_text_embedding=None,
-    ):
-        super().__init__(
-            vae,
-            text_encoder,
-            tokenizer,
-            unet,
-            controlnet,
-            scheduler,
-            safety_checker,
-            feature_extractor,
-            image_encoder,
-            requires_safety_checker,
-                )
-
-        self.register_modules(
-            dino_controlnet=dino_controlnet,
-        )
-
-        self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.dino_image_processor = lambda x: x / 127.5 -1.
-
-        self.default_denoising_steps = default_denoising_steps
-        self.default_processing_resolution = default_processing_resolution
-        self.prompt = prompt
-        self.prompt_embeds = None
-        self.empty_text_embedding = empty_text_embedding
-        self.prior = DINOv2_Encoder(size=672)
-
-    def check_inputs(
-        self,
-        image: PipelineImageInput,
-        num_inference_steps: int,
-        ensemble_size: int,
-        processing_resolution: int,
-        resample_method_input: str,
-        resample_method_output: str,
-        batch_size: int,
-        ensembling_kwargs: Optional[Dict[str, Any]],
-        latents: Optional[torch.Tensor],
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
-        output_type: str,
-        output_uncertainty: bool,
-    ) -> int:
-        if num_inference_steps is None:
-            raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
-        if num_inference_steps < 1:
-            raise ValueError("`num_inference_steps` must be positive.")
-        if ensemble_size < 1:
-            raise ValueError("`ensemble_size` must be positive.")
-        if ensemble_size == 2:
-            logger.warning(
-                "`ensemble_size` == 2 results are similar to no ensembling (1); "
-                "consider increasing the value to at least 3."
-            )
-        if ensemble_size == 1 and output_uncertainty:
-            raise ValueError(
-                "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
-                "greater than 1."
-            )
-        if processing_resolution is None:
-            raise ValueError(
-                "`processing_resolution` is not specified and could not be resolved from the model config."
-            )
-        if processing_resolution < 0:
-            raise ValueError(
-                "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
-                "downsampled processing."
-            )
-        if processing_resolution % self.vae_scale_factor != 0:
-            raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
-        if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
-            raise ValueError(
-                "`resample_method_input` takes string values compatible with PIL library: "
-                "nearest, nearest-exact, bilinear, bicubic, area."
-            )
-        if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
-            raise ValueError(
-                "`resample_method_output` takes string values compatible with PIL library: "
-                "nearest, nearest-exact, bilinear, bicubic, area."
-            )
-        if batch_size < 1:
-            raise ValueError("`batch_size` must be positive.")
-        if output_type not in ["pt", "np"]:
-            raise ValueError("`output_type` must be one of `pt` or `np`.")
-        if latents is not None and generator is not None:
-            raise ValueError("`latents` and `generator` cannot be used together.")
-        if ensembling_kwargs is not None:
-            if not isinstance(ensembling_kwargs, dict):
-                raise ValueError("`ensembling_kwargs` must be a dictionary.")
-            if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"):
-                raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.")
-
-        # image checks
-        num_images = 0
-        W, H = None, None
-        if not isinstance(image, list):
-            image = [image]
-        for i, img in enumerate(image):
-            if isinstance(img, np.ndarray) or torch.is_tensor(img):
-                if img.ndim not in (2, 3, 4):
-                    raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
-                H_i, W_i = img.shape[-2:]
-                N_i = 1
-                if img.ndim == 4:
-                    N_i = img.shape[0]
-            elif isinstance(img, Image.Image):
-                W_i, H_i = img.size
-                N_i = 1
-            else:
-                raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
-            if W is None:
-                W, H = W_i, H_i
-            elif (W, H) != (W_i, H_i):
-                raise ValueError(
-                    f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
-                )
-            num_images += N_i
-
-        # latents checks
-        if latents is not None:
-            if not torch.is_tensor(latents):
-                raise ValueError("`latents` must be a torch.Tensor.")
-            if latents.dim() != 4:
-                raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")
-
-            if processing_resolution > 0:
-                max_orig = max(H, W)
-                new_H = H * processing_resolution // max_orig
-                new_W = W * processing_resolution // max_orig
-                if new_H == 0 or new_W == 0:
-                    raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
-                W, H = new_W, new_H
-            w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
-            h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
-            shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w)
-
-            if latents.shape != shape_expected:
-                raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")
-
-        # generator checks
-        if generator is not None:
-            if isinstance(generator, list):
-                if len(generator) != num_images * ensemble_size:
-                    raise ValueError(
-                        "The number of generators must match the total number of ensemble members for all input images."
-                    )
-                if not all(g.device.type == generator[0].device.type for g in generator):
-                    raise ValueError("`generator` device placement is not consistent in the list.")
-            elif not isinstance(generator, torch.Generator):
-                raise ValueError(f"Unsupported generator type: {type(generator)}.")
-
-        return num_images
-
-    def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
-        if not hasattr(self, "_progress_bar_config"):
-            self._progress_bar_config = {}
-        elif not isinstance(self._progress_bar_config, dict):
-            raise ValueError(
-                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
-            )
-
-        progress_bar_config = dict(**self._progress_bar_config)
-        progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
-        progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
-        if iterable is not None:
-            return tqdm(iterable, **progress_bar_config)
-        elif total is not None:
-            return tqdm(total=total, **progress_bar_config)
-        else:
-            raise ValueError("Either `total` or `iterable` has to be defined.")
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        image: PipelineImageInput,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_inference_steps: Optional[int] = None,
-        ensemble_size: int = 1,
-        processing_resolution: Optional[int] = None,
-        match_input_resolution: bool = True,
-        resample_method_input: str = "bilinear",
-        resample_method_output: str = "bilinear",
-        batch_size: int = 1,
-        ensembling_kwargs: Optional[Dict[str, Any]] = None,
-        latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
-        output_type: str = "np",
-        output_uncertainty: bool = False,
-        output_latent: bool = False,
-        return_dict: bool = True,
-    ):
-        """
-        Function invoked when calling the pipeline.
-
-        Args:
-            image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
-                `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For
-                arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
-                by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
-                three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
-                same width and height.
-            num_inference_steps (`int`, *optional*, defaults to `None`):
-                Number of denoising diffusion steps during inference. The default value `None` results in automatic
-                selection. The number of steps should be at least 10 with the full Marigold models, and between 1 and 4
-                for Marigold-LCM models.
-            ensemble_size (`int`, defaults to `1`):
-                Number of ensemble predictions. Recommended values are 5 and higher for better precision, or 1 for
-                faster inference.
-            processing_resolution (`int`, *optional*, defaults to `None`):
-                Effective processing resolution. When set to `0`, matches the larger input image dimension. This
-                produces crisper predictions, but may also lead to the overall loss of global context. The default
-                value `None` resolves to the optimal value from the model config.
-            match_input_resolution (`bool`, *optional*, defaults to `True`):
-                When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
-                side of the output will equal to `processing_resolution`.
-            resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
-                Resampling method used to resize input images to `processing_resolution`. The accepted values are:
-                `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
-            resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
-                Resampling method used to resize output predictions to match the input resolution. The accepted values
-                are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
-            batch_size (`int`, *optional*, defaults to `1`):
-                Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
-            ensembling_kwargs (`dict`, *optional*, defaults to `None`)
-                Extra dictionary with arguments for precise ensembling control. The following options are available:
-                - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in
-                  every pixel location, can be either `"closest"` or `"mean"`.
-            latents (`torch.Tensor`, *optional*, defaults to `None`):
-                Latent noise tensors to replace the random initialization. These can be taken from the previous
-                function call's output.
-            generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
-                Random number generator object to ensure reproducibility.
-            output_type (`str`, *optional*, defaults to `"np"`):
-                Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
-                values are: `"np"` (numpy array) or `"pt"` (torch tensor).
-            output_uncertainty (`bool`, *optional*, defaults to `False`):
-                When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
-                the `ensemble_size` argument is set to a value above 2.
-            output_latent (`bool`, *optional*, defaults to `False`):
-                When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
-                within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
-                `latents` argument.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a
-                `tuple` is returned where the first element is the prediction, the second element is the uncertainty
-                (or `None`), and the third is the latent (or `None`).
-        """
-
-        # 0. Resolving variables.
-        device = self._execution_device
-        dtype = self.dtype
-
-        # Model-specific optimal default values leading to fast and reasonable results.
-        if num_inference_steps is None:
-            num_inference_steps = self.default_denoising_steps
-        if processing_resolution is None:
-            processing_resolution = self.default_processing_resolution
-
-
-        image, padding, original_resolution = self.image_processor.preprocess(
-            image, processing_resolution, resample_method_input, device, dtype
-        )  # [N,3,PPH,PPW]
-
-        image_latent, gaus_noise = self.prepare_latents(
-            image, latents, generator, ensemble_size, batch_size
-        )  # [N,4,h,w], [N,4,h,w]
-
-        # 0. X_start latent obtain
-        predictor = self.x_start_pipeline(image, latents=gaus_noise, 
-                                          processing_resolution=processing_resolution, skip_preprocess=True)
-        x_start_latent = predictor.latent
-
-        # 1. Check inputs.
-        num_images = self.check_inputs(
-            image,
-            num_inference_steps,
-            ensemble_size,
-            processing_resolution,
-            resample_method_input,
-            resample_method_output,
-            batch_size,
-            ensembling_kwargs,
-            latents,
-            generator,
-            output_type,
-            output_uncertainty,
-        )
-
-
-        # 2. Prepare empty text conditioning.
-        # Model invocation: self.tokenizer, self.text_encoder.
-        if self.empty_text_embedding is None:
-            prompt = ""
-            text_inputs = self.tokenizer(
-                prompt,
-                padding="do_not_pad",
-                max_length=self.tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-            text_input_ids = text_inputs.input_ids.to(device)
-            self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]
-
-
-
-        # 3. prepare prompt
-        if self.prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                self.prompt,
-                device,
-                num_images_per_prompt,
-                False,
-                negative_prompt,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=None,
-                lora_scale=None,
-                clip_skip=None,
-            )
-            self.prompt_embeds = prompt_embeds
-            self.negative_prompt_embeds = negative_prompt_embeds
-
-
-
-        # 5. dino guider features obtaining
-        ## TODO different case-1
-        dino_features = self.prior(image)
-        dino_features = self.dino_controlnet.dino_controlnet_cond_embedding(dino_features)
-        dino_features = self.match_noisy(dino_features, x_start_latent)
-
-        del (
-                image,
-        )
-
-        # 7. denoise sampling, using heuritic sampling proposed by Ye.
-
-        t_start = self.x_start_pipeline.t_start
-        self.scheduler.set_timesteps(num_inference_steps, t_start=t_start,device=device)
-
-        cond_scale =controlnet_conditioning_scale
-        pred_latent = x_start_latent
-
-        cur_step = 0
-
-        # dino controlnet
-        dino_down_block_res_samples, dino_mid_block_res_sample = self.dino_controlnet(
-            dino_features.detach(),
-            0, # not depend on time steps
-            encoder_hidden_states=self.prompt_embeds,
-            conditioning_scale=cond_scale,
-            guess_mode=False,
-            return_dict=False,
-        )
-        assert dino_mid_block_res_sample == None
-
-        pred_latents = []
-
-        last_pred_latent = pred_latent
-        for (t, prev_t) in self.progress_bar(zip(self.scheduler.timesteps,self.scheduler.prev_timesteps), leave=False, desc="Diffusion steps..."):
-
-            _dino_down_block_res_samples = [dino_down_block_res_sample for dino_down_block_res_sample in dino_down_block_res_samples]  # copy, avoid repeat quiery
-
-            # controlnet
-            down_block_res_samples, mid_block_res_sample = self.controlnet(
-                image_latent.detach(),
-                t,
-                encoder_hidden_states=self.prompt_embeds,
-                conditioning_scale=cond_scale,
-                guess_mode=False,
-                return_dict=False,
-            )
-
-            # SG-DRN
-            noise = self.dino_unet_forward(
-                self.unet,
-                pred_latent,
-                t,
-                encoder_hidden_states=self.prompt_embeds,
-                down_block_additional_residuals=down_block_res_samples,
-                mid_block_additional_residual=mid_block_res_sample,
-                dino_down_block_additional_residuals= _dino_down_block_res_samples,
-                return_dict=False,
-            )[0]  # [B,4,h,w]
-
-            pred_latents.append(noise)
-            # ddim steps
-            out = self.scheduler.step(
-                noise, t, prev_t, pred_latent, gaus_noise = gaus_noise, generator=generator, cur_step=cur_step+1  # NOTE that cur_step dirs to next_step
-            )# [B,4,h,w]
-            pred_latent = out.prev_sample
-
-            cur_step += 1
-
-        del (
-            image_latent,
-            dino_features,
-        )
-        pred_latent = pred_latents[-1]  # using x0
-
-        # decoder
-        prediction = self.decode_prediction(pred_latent)
-        prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,3,PH,PW]
-        prediction = self.image_processor.resize_antialias(prediction, original_resolution, resample_method_output, is_aa=False)  # [N,3,H,W]
-
-        if match_input_resolution:
-            prediction = self.image_processor.resize_antialias(
-                prediction, original_resolution, resample_method_output, is_aa=False
-            )  # [N,3,H,W]
-
-        if match_input_resolution:
-            prediction = self.image_processor.resize_antialias(
-                prediction, original_resolution, resample_method_output, is_aa=False
-            )  # [N,3,H,W]
-        prediction = self.normalize_normals(prediction)  # [N,3,H,W]
-
-        if output_type == "np":
-            prediction = self.image_processor.pt_to_numpy(prediction)  # [N,H,W,3]
-            prediction = prediction.clip(min=-1, max=1)
-
-        # 11. Offload all models
-        self.maybe_free_model_hooks()
-
-        return StableNormalOutput(
-            prediction=prediction,
-            latent=pred_latent,
-            gaus_noise=gaus_noise
-        )
-
-    # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents
-    def prepare_latents(
-        self,
-        image: torch.Tensor,
-        latents: Optional[torch.Tensor],
-        generator: Optional[torch.Generator],
-        ensemble_size: int,
-        batch_size: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        def retrieve_latents(encoder_output):
-            if hasattr(encoder_output, "latent_dist"):
-                return encoder_output.latent_dist.mode()
-            elif hasattr(encoder_output, "latents"):
-                return encoder_output.latents
-            else:
-                raise AttributeError("Could not access latents of provided encoder_output")
-
-
-
-        image_latent = torch.cat(
-            [
-                retrieve_latents(self.vae.encode(image[i : i + batch_size]))
-                for i in range(0, image.shape[0], batch_size)
-            ],
-            dim=0,
-        )  # [N,4,h,w]
-        image_latent = image_latent * self.vae.config.scaling_factor
-        image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]
-
-        pred_latent = latents
-        if pred_latent is None:
-
-
-            pred_latent = randn_tensor(
-                image_latent.shape,
-                generator=generator,
-                device=image_latent.device,
-                dtype=image_latent.dtype,
-            )  # [N*E,4,h,w]
-
-        return image_latent, pred_latent
-
-    def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor:
-        if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
-            raise ValueError(
-                f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
-            )
-
-        prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0]  # [B,3,H,W]
-
-        return prediction  # [B,3,H,W]
-
-    @staticmethod
-    def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-        if normals.dim() != 4 or normals.shape[1] != 3:
-            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
-
-        norm = torch.norm(normals, dim=1, keepdim=True)
-        normals /= norm.clamp(min=eps)
-
-        return normals
-
-    @staticmethod
-    def match_noisy(dino, noisy):
-        _, __, dino_h, dino_w =  dino.shape
-        _, __, h, w =  noisy.shape
-
-        if h == dino_h and w == dino_w:
-            return dino
-        else:
-            return F.interpolate(dino, (h, w), mode='bilinear')
-
-
-
-
-
-
-
-
-
-
-    @staticmethod
-    def dino_unet_forward(
-        self,  # NOTE that repurpose to UNet
-        sample: torch.Tensor,
-        timestep: Union[torch.Tensor, float, int],
-        encoder_hidden_states: torch.Tensor,
-        class_labels: Optional[torch.Tensor] = None,
-        timestep_cond: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
-        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
-        mid_block_additional_residual: Optional[torch.Tensor] = None,
-        dino_down_block_additional_residuals: Optional[torch.Tensor] = None,
-        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        return_dict: bool = True,
-    ) -> Union[UNet2DConditionOutput, Tuple]:
-        r"""
-        The [`UNet2DConditionModel`] forward method.
-
-        Args:
-            sample (`torch.Tensor`):
-                The noisy input tensor with the following shape `(batch, channel, height, width)`.
-            timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input.
-            encoder_hidden_states (`torch.Tensor`):
-                The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
-            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
-                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
-            timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
-                Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
-                through the `self.time_embedding` layer to obtain the timestep embeddings.
-            attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
-                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
-                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
-                negative values to the attention scores corresponding to "discard" tokens.
-            cross_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            added_cond_kwargs: (`dict`, *optional*):
-                A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
-                are passed along to the UNet blocks.
-            down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
-                A tuple of tensors that if specified are added to the residuals of down unet blocks.
-            mid_block_additional_residual: (`torch.Tensor`, *optional*):
-                A tensor that if specified is added to the residual of the middle unet block.
-            down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
-                additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
-            encoder_attention_mask (`torch.Tensor`):
-                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
-                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
-                which adds large negative values to the attention scores corresponding to "discard" tokens.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
-                tuple.
-
-        Returns:
-            [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
-                If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned,
-                otherwise a `tuple` is returned where the first element is the sample tensor.
-        """
-        # By default samples have to be AT least a multiple of the overall upsampling factor.
-        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
-        # However, the upsampling interpolation output size can be forced to fit any upsampling size
-        # on the fly if necessary.
-
-
-        default_overall_up_factor = 2**self.num_upsamplers
-
-        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-        forward_upsample_size = False
-        upsample_size = None
-
-        for dim in sample.shape[-2:]:
-            if dim % default_overall_up_factor != 0:
-                # Forward upsample size to force interpolation output size.
-                forward_upsample_size = True
-                break
-
-        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
-        # expects mask of shape:
-        #   [batch, key_tokens]
-        # adds singleton query_tokens dimension:
-        #   [batch,                    1, key_tokens]
-        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
-        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
-        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
-        if attention_mask is not None:
-            # assume that mask is expressed as:
-            #   (1 = keep,      0 = discard)
-            # convert mask into a bias that can be added to attention scores:
-            #       (keep = +0,     discard = -10000.0)
-            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
-            attention_mask = attention_mask.unsqueeze(1)
-
-        # convert encoder_attention_mask to a bias the same way we do for attention_mask
-        if encoder_attention_mask is not None:
-            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
-            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
-
-        # 0. center input if necessary
-        if self.config.center_input_sample:
-            sample = 2 * sample - 1.0
-
-        # 1. time
-        t_emb = self.get_time_embed(sample=sample, timestep=timestep)
-        emb = self.time_embedding(t_emb, timestep_cond)
-        aug_emb = None
-
-        class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
-        if class_emb is not None:
-            if self.config.class_embeddings_concat:
-                emb = torch.cat([emb, class_emb], dim=-1)
-            else:
-                emb = emb + class_emb
-
-        aug_emb = self.get_aug_embed(
-            emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
-        )
-        if self.config.addition_embed_type == "image_hint":
-            aug_emb, hint = aug_emb
-            sample = torch.cat([sample, hint], dim=1)
-
-        emb = emb + aug_emb if aug_emb is not None else emb
-
-        if self.time_embed_act is not None:
-            emb = self.time_embed_act(emb)
-
-        encoder_hidden_states = self.process_encoder_hidden_states(
-            encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
-        )
-
-        # 2. pre-process
-        sample = self.conv_in(sample)
-
-        # 2.5 GLIGEN position net
-        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
-            cross_attention_kwargs = cross_attention_kwargs.copy()
-            gligen_args = cross_attention_kwargs.pop("gligen")
-            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
-
-        # 3. down
-        # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
-        # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
-        if cross_attention_kwargs is not None:
-            cross_attention_kwargs = cross_attention_kwargs.copy()
-            lora_scale = cross_attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-
-        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
-        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
-        is_adapter = down_intrablock_additional_residuals is not None
-        # maintain backward compatibility for legacy usage, where
-        #       T2I-Adapter and ControlNet both use down_block_additional_residuals arg
-        #       but can only use one or the other
-        if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
-            deprecate(
-                "T2I should not use down_block_additional_residuals",
-                "1.3.0",
-                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
-                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
-                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
-                standard_warn=False,
-            )
-            down_intrablock_additional_residuals = down_block_additional_residuals
-            is_adapter = True
-
-
-
-        def residual_downforward(
-            self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None,
-            additional_residuals: Optional[torch.Tensor] = None,
-            *args, **kwargs,
-        ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
-            if len(args) > 0 or kwargs.get("scale", None) is not None:
-                deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
-                deprecate("scale", "1.0.0", deprecation_message)
-
-            output_states = ()
-
-            for resnet in self.resnets:
-                if self.training and self.gradient_checkpointing:
-
-                    def create_custom_forward(module):
-                        def custom_forward(*inputs):
-                            return module(*inputs)
-
-                        return custom_forward
-
-                    if is_torch_version(">=", "1.11.0"):
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                        )
-                    else:
-                        hidden_states = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(resnet), hidden_states, temb
-                        )
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-                    hidden_states += additional_residuals.pop(0)
-
-
-                output_states = output_states + (hidden_states,)
-
-            if self.downsamplers is not None:
-                for downsampler in self.downsamplers:
-                    hidden_states = downsampler(hidden_states)
-                    hidden_states += additional_residuals.pop(0)
-
-                output_states = output_states + (hidden_states,)
-
-            return hidden_states, output_states
-
-
-        def residual_blockforward(
-            self,  ## NOTE that repurpose to unet_blocks
-            hidden_states: torch.Tensor,
-            temb: Optional[torch.Tensor] = None,
-            encoder_hidden_states: Optional[torch.Tensor] = None,
-            attention_mask: Optional[torch.Tensor] = None,
-            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-            encoder_attention_mask: Optional[torch.Tensor] = None,
-            additional_residuals: Optional[torch.Tensor] = None,
-        ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
-            if cross_attention_kwargs is not None:
-                if cross_attention_kwargs.get("scale", None) is not None:
-                    logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
-
-
-
-            output_states = ()
-
-            blocks = list(zip(self.resnets, self.attentions))
-
-            for i, (resnet, attn) in enumerate(blocks):
-                if self.training and self.gradient_checkpointing:
-
-                    def create_custom_forward(module, return_dict=None):
-                        def custom_forward(*inputs):
-                            if return_dict is not None:
-                                return module(*inputs, return_dict=return_dict)
-                            else:
-                                return module(*inputs)
-
-                        return custom_forward
-
-                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(resnet),
-                        hidden_states,
-                        temb,
-                        **ckpt_kwargs,
-                    )
-                    hidden_states = attn(
-                        hidden_states,
-                        encoder_hidden_states=encoder_hidden_states,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        attention_mask=attention_mask,
-                        encoder_attention_mask=encoder_attention_mask,
-                        return_dict=False,
-                    )[0]
-                else:
-                    hidden_states = resnet(hidden_states, temb)
-                    hidden_states = attn(
-                        hidden_states,
-                        encoder_hidden_states=encoder_hidden_states,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        attention_mask=attention_mask,
-                        encoder_attention_mask=encoder_attention_mask,
-                        return_dict=False,
-                    )[0]
-
-                hidden_states += additional_residuals.pop(0)
-
-                output_states = output_states + (hidden_states,)
-
-            if self.downsamplers is not None:
-                for downsampler in self.downsamplers:
-                    hidden_states = downsampler(hidden_states)
-                    hidden_states += additional_residuals.pop(0)
-
-                output_states = output_states + (hidden_states,)
-
-            return hidden_states, output_states
-
-
-        down_intrablock_additional_residuals = dino_down_block_additional_residuals
-
-        sample += down_intrablock_additional_residuals.pop(0)
-        down_block_res_samples = (sample,)
-
-        for downsample_block in self.down_blocks:
-            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
-
-                sample, res_samples = residual_blockforward(
-                    downsample_block,
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    encoder_attention_mask=encoder_attention_mask,
-                    additional_residuals = down_intrablock_additional_residuals,
-                )
-
-            else:
-                sample, res_samples = residual_downforward(
-                    downsample_block,
-                    hidden_states=sample,
-                    temb=emb,
-                    additional_residuals = down_intrablock_additional_residuals,
-                        )
-
-
-            down_block_res_samples += res_samples
-
-
-        if is_controlnet:
-            new_down_block_res_samples = ()
-
-            for down_block_res_sample, down_block_additional_residual in zip(
-                down_block_res_samples, down_block_additional_residuals
-            ):
-                down_block_res_sample = down_block_res_sample + down_block_additional_residual
-                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
-
-            down_block_res_samples = new_down_block_res_samples
-
-        # 4. mid
-        if self.mid_block is not None:
-            if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
-                sample = self.mid_block(
-                    sample,
-                    emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    encoder_attention_mask=encoder_attention_mask,
-                )
-            else:
-                sample = self.mid_block(sample, emb)
-
-            # To support T2I-Adapter-XL
-            if (
-                is_adapter
-                and len(down_intrablock_additional_residuals) > 0
-                and sample.shape == down_intrablock_additional_residuals[0].shape
-            ):
-                sample += down_intrablock_additional_residuals.pop(0)
-
-        if is_controlnet:
-            sample = sample + mid_block_additional_residual
-
-        # 5. up
-        for i, upsample_block in enumerate(self.up_blocks):
-            is_final_block = i == len(self.up_blocks) - 1
-
-            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
-
-            # if we have not reached the final block and need to forward the
-            # upsample size, we do it here
-            if not is_final_block and forward_upsample_size:
-                upsample_size = down_block_res_samples[-1].shape[2:]
-
-            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    upsample_size=upsample_size,
-                    attention_mask=attention_mask,
-                    encoder_attention_mask=encoder_attention_mask,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    upsample_size=upsample_size,
-                )
-
-        # 6. post-process
-        if self.conv_norm_out:
-            sample = self.conv_norm_out(sample)
-            sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not return_dict:
-            return (sample,)
-
-        return UNet2DConditionOutput(sample=sample)
-
-
-
-    @staticmethod
-    def ensemble_normals(
-        normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest"
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """
-        Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
-        the number of ensemble members for a given prediction of size `(H x W)`.
-
-        Args:
-            normals (`torch.Tensor`):
-                Input ensemble normals maps.
-            output_uncertainty (`bool`, *optional*, defaults to `False`):
-                Whether to output uncertainty map.
-            reduction (`str`, *optional*, defaults to `"closest"`):
-                Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
-                `"mean"`.
-
-        Returns:
-            A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
-            uncertainties of shape `(1, 1, H, W)`.
-        """
-        if normals.dim() != 4 or normals.shape[1] != 3:
-            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
-        if reduction not in ("closest", "mean"):
-            raise ValueError(f"Unrecognized reduction method: {reduction}.")
-
-        mean_normals = normals.mean(dim=0, keepdim=True)  # [1,3,H,W]
-        mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals)  # [1,3,H,W]
-
-        sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True)  # [E,1,H,W]
-        sim_cos = sim_cos.clamp(-1, 1)  # required to avoid NaN in uncertainty with fp16
-
-        uncertainty = None
-        if output_uncertainty:
-            uncertainty = sim_cos.arccos()  # [E,1,H,W]
-            uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi  # [1,1,H,W]
-
-        if reduction == "mean":
-            return mean_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
-
-        closest_indices = sim_cos.argmax(dim=0, keepdim=True)  # [1,1,H,W]
-        closest_indices = closest_indices.repeat(1, 3, 1, 1)  # [1,3,H,W]
-        closest_normals = torch.gather(normals, 0, closest_indices)  # [1,3,H,W]
-
-        return closest_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps

stablenormal/pipeline_yoso_normal.py

@@ -1,727 +0,0 @@
-# Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved.
-# Copyright 2024 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.
-# --------------------------------------------------------------------------
-# More information and citation instructions are available on the
-# --------------------------------------------------------------------------
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from PIL import Image
-from tqdm.auto import tqdm
-from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
-
-
-from diffusers.image_processor import PipelineImageInput
-from diffusers.models import (
-    AutoencoderKL,
-    UNet2DConditionModel,
-	ControlNetModel,
-)
-from diffusers.schedulers import (
-	DDIMScheduler
-)
-
-from diffusers.utils import (
-    BaseOutput,
-    logging,
-    replace_example_docstring,
-)
-
-
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.controlnet import StableDiffusionControlNetPipeline
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.marigold.marigold_image_processing import MarigoldImageProcessor
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-
-import pdb
-
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
-Examples:
-```py
->>> import diffusers
->>> import torch
-
->>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
-...     "prs-eth/marigold-normals-lcm-v0-1", variant="fp16", torch_dtype=torch.float16
-... ).to("cuda")
-
->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
->>> normals = pipe(image)
-
->>> vis = pipe.image_processor.visualize_normals(normals.prediction)
->>> vis[0].save("einstein_normals.png")
-```
-"""
-
-
-@dataclass
-class YosoNormalsOutput(BaseOutput):
-    """
-    Output class for Marigold monocular normals prediction pipeline.
-
-    Args:
-        prediction (`np.ndarray`, `torch.Tensor`):
-            Predicted normals with values in the range [-1, 1]. The shape is always $numimages \times 3 \times height
-            \times width$, regardless of whether the images were passed as a 4D array or a list.
-        uncertainty (`None`, `np.ndarray`, `torch.Tensor`):
-            Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages
-            \times 1 \times height \times width$.
-        latent (`None`, `torch.Tensor`):
-            Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
-            The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$.
-    """
-
-    prediction: Union[np.ndarray, torch.Tensor]
-    latent: Union[None, torch.Tensor]
-    gaus_noise: Union[None, torch.Tensor]
-
-
-class YOSONormalsPipeline(StableDiffusionControlNetPipeline):
-    """ Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io.
-    Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    The pipeline also inherits the following loading methods:
-        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
-        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
-        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
-        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
-        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
-
-    Args:
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
-        text_encoder ([`~transformers.CLIPTextModel`]):
-            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
-        tokenizer ([`~transformers.CLIPTokenizer`]):
-            A `CLIPTokenizer` to tokenize text.
-        unet ([`UNet2DConditionModel`]):
-            A `UNet2DConditionModel` to denoise the encoded image latents.
-        controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
-            Provides additional conditioning to the `unet` during the denoising process. If you set multiple
-            ControlNets as a list, the outputs from each ControlNet are added together to create one combined
-            additional conditioning.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
-            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
-        safety_checker ([`StableDiffusionSafetyChecker`]):
-            Classification module that estimates whether generated images could be considered offensive or harmful.
-            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
-            about a model's potential harms.
-        feature_extractor ([`~transformers.CLIPImageProcessor`]):
-            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
-    """
-
-    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
-    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
-    _exclude_from_cpu_offload = ["safety_checker"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        tokenizer: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel]],
-        scheduler: Union[DDIMScheduler],
-        safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
-        image_encoder: CLIPVisionModelWithProjection = None,
-        requires_safety_checker: bool = True,
-        default_denoising_steps: Optional[int] = 1,
-		default_processing_resolution: Optional[int] = 768,
-        prompt="",
-        empty_text_embedding=None,
-        t_start: Optional[int] = 401,
-    ):
-        super().__init__(
-            vae,
-            text_encoder,
-            tokenizer,
-            unet,
-            controlnet,
-            scheduler,
-            safety_checker,
-            feature_extractor,
-            image_encoder,
-            requires_safety_checker,
-                )
-
-        # TODO yoso ImageProcessor
-        self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.control_image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.default_denoising_steps = default_denoising_steps
-        self.default_processing_resolution = default_processing_resolution
-        self.prompt = prompt
-        self.prompt_embeds = None
-        self.empty_text_embedding = empty_text_embedding
-        self.t_start= t_start # target_out latents
-
-    def check_inputs(
-        self,
-        image: PipelineImageInput,
-        num_inference_steps: int,
-        ensemble_size: int,
-        processing_resolution: int,
-        resample_method_input: str,
-        resample_method_output: str,
-        batch_size: int,
-        ensembling_kwargs: Optional[Dict[str, Any]],
-        latents: Optional[torch.Tensor],
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
-        output_type: str,
-        output_uncertainty: bool,
-    ) -> int:
-        if num_inference_steps is None:
-            raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
-        if num_inference_steps < 1:
-            raise ValueError("`num_inference_steps` must be positive.")
-        if ensemble_size < 1:
-            raise ValueError("`ensemble_size` must be positive.")
-        if ensemble_size == 2:
-            logger.warning(
-                "`ensemble_size` == 2 results are similar to no ensembling (1); "
-                "consider increasing the value to at least 3."
-            )
-        if ensemble_size == 1 and output_uncertainty:
-            raise ValueError(
-                "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
-                "greater than 1."
-            )
-        if processing_resolution is None:
-            raise ValueError(
-                "`processing_resolution` is not specified and could not be resolved from the model config."
-            )
-        if processing_resolution < 0:
-            raise ValueError(
-                "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
-                "downsampled processing."
-            )
-        if processing_resolution % self.vae_scale_factor != 0:
-            raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
-        if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
-            raise ValueError(
-                "`resample_method_input` takes string values compatible with PIL library: "
-                "nearest, nearest-exact, bilinear, bicubic, area."
-            )
-        if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
-            raise ValueError(
-                "`resample_method_output` takes string values compatible with PIL library: "
-                "nearest, nearest-exact, bilinear, bicubic, area."
-            )
-        if batch_size < 1:
-            raise ValueError("`batch_size` must be positive.")
-        if output_type not in ["pt", "np"]:
-            raise ValueError("`output_type` must be one of `pt` or `np`.")
-        if latents is not None and generator is not None:
-            raise ValueError("`latents` and `generator` cannot be used together.")
-        if ensembling_kwargs is not None:
-            if not isinstance(ensembling_kwargs, dict):
-                raise ValueError("`ensembling_kwargs` must be a dictionary.")
-            if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"):
-                raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.")
-
-        # image checks
-        num_images = 0
-        W, H = None, None
-        if not isinstance(image, list):
-            image = [image]
-        for i, img in enumerate(image):
-            if isinstance(img, np.ndarray) or torch.is_tensor(img):
-                if img.ndim not in (2, 3, 4):
-                    raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
-                H_i, W_i = img.shape[-2:]
-                N_i = 1
-                if img.ndim == 4:
-                    N_i = img.shape[0]
-            elif isinstance(img, Image.Image):
-                W_i, H_i = img.size
-                N_i = 1
-            else:
-                raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
-            if W is None:
-                W, H = W_i, H_i
-            elif (W, H) != (W_i, H_i):
-                raise ValueError(
-                    f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
-                )
-            num_images += N_i
-
-        # latents checks
-        if latents is not None:
-            if not torch.is_tensor(latents):
-                raise ValueError("`latents` must be a torch.Tensor.")
-            if latents.dim() != 4:
-                raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")
-
-            if processing_resolution > 0:
-                max_orig = max(H, W)
-                new_H = H * processing_resolution // max_orig
-                new_W = W * processing_resolution // max_orig
-                if new_H == 0 or new_W == 0:
-                    raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
-                W, H = new_W, new_H
-            w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
-            h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
-            shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w)
-
-            if latents.shape != shape_expected:
-                raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")
-
-        # generator checks
-        if generator is not None:
-            if isinstance(generator, list):
-                if len(generator) != num_images * ensemble_size:
-                    raise ValueError(
-                        "The number of generators must match the total number of ensemble members for all input images."
-                    )
-                if not all(g.device.type == generator[0].device.type for g in generator):
-                    raise ValueError("`generator` device placement is not consistent in the list.")
-            elif not isinstance(generator, torch.Generator):
-                raise ValueError(f"Unsupported generator type: {type(generator)}.")
-
-        return num_images
-
-    def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
-        if not hasattr(self, "_progress_bar_config"):
-            self._progress_bar_config = {}
-        elif not isinstance(self._progress_bar_config, dict):
-            raise ValueError(
-                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
-            )
-
-        progress_bar_config = dict(**self._progress_bar_config)
-        progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
-        progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
-        if iterable is not None:
-            return tqdm(iterable, **progress_bar_config)
-        elif total is not None:
-            return tqdm(total=total, **progress_bar_config)
-        else:
-            raise ValueError("Either `total` or `iterable` has to be defined.")
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        image: PipelineImageInput,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_inference_steps: Optional[int] = None,
-        ensemble_size: int = 1,
-        processing_resolution: Optional[int] = None,
-        match_input_resolution: bool = True,
-        resample_method_input: str = "bilinear",
-        resample_method_output: str = "bilinear",
-        batch_size: int = 1,
-        ensembling_kwargs: Optional[Dict[str, Any]] = None,
-        latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
-        output_type: str = "np",
-        output_uncertainty: bool = False,
-        output_latent: bool = False,
-        skip_preprocess: bool = False,
-        return_dict: bool = True,
-        **kwargs,
-    ):
-        """
-        Function invoked when calling the pipeline.
-
-        Args:
-            image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
-                `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For
-                arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
-                by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
-                three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
-                same width and height.
-            num_inference_steps (`int`, *optional*, defaults to `None`):
-                Number of denoising diffusion steps during inference. The default value `None` results in automatic
-                selection. The number of steps should be at least 10 with the full Marigold models, and between 1 and 4
-                for Marigold-LCM models.
-            ensemble_size (`int`, defaults to `1`):
-                Number of ensemble predictions. Recommended values are 5 and higher for better precision, or 1 for
-                faster inference.
-            processing_resolution (`int`, *optional*, defaults to `None`):
-                Effective processing resolution. When set to `0`, matches the larger input image dimension. This
-                produces crisper predictions, but may also lead to the overall loss of global context. The default
-                value `None` resolves to the optimal value from the model config.
-            match_input_resolution (`bool`, *optional*, defaults to `True`):
-                When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
-                side of the output will equal to `processing_resolution`.
-            resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
-                Resampling method used to resize input images to `processing_resolution`. The accepted values are:
-                `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
-            resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
-                Resampling method used to resize output predictions to match the input resolution. The accepted values
-                are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
-            batch_size (`int`, *optional*, defaults to `1`):
-                Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
-            ensembling_kwargs (`dict`, *optional*, defaults to `None`)
-                Extra dictionary with arguments for precise ensembling control. The following options are available:
-                - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in
-                  every pixel location, can be either `"closest"` or `"mean"`.
-            latents (`torch.Tensor`, *optional*, defaults to `None`):
-                Latent noise tensors to replace the random initialization. These can be taken from the previous
-                function call's output.
-            generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
-                Random number generator object to ensure reproducibility.
-            output_type (`str`, *optional*, defaults to `"np"`):
-                Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
-                values are: `"np"` (numpy array) or `"pt"` (torch tensor).
-            output_uncertainty (`bool`, *optional*, defaults to `False`):
-                When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
-                the `ensemble_size` argument is set to a value above 2.
-            output_latent (`bool`, *optional*, defaults to `False`):
-                When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
-                within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
-                `latents` argument.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a
-                `tuple` is returned where the first element is the prediction, the second element is the uncertainty
-                (or `None`), and the third is the latent (or `None`).
-        """
-
-        # 0. Resolving variables.
-        device = self._execution_device
-        dtype = self.dtype
-
-        # Model-specific optimal default values leading to fast and reasonable results.
-        if num_inference_steps is None:
-            num_inference_steps = self.default_denoising_steps
-        if processing_resolution is None:
-            processing_resolution = self.default_processing_resolution
-
-        # 1. Check inputs.
-        num_images = self.check_inputs(
-            image,
-            num_inference_steps,
-            ensemble_size,
-            processing_resolution,
-            resample_method_input,
-            resample_method_output,
-            batch_size,
-            ensembling_kwargs,
-            latents,
-            generator,
-            output_type,
-            output_uncertainty,
-        )
-
-
-        # 2. Prepare empty text conditioning.
-        # Model invocation: self.tokenizer, self.text_encoder.
-        if self.empty_text_embedding is None:
-            prompt = ""
-            text_inputs = self.tokenizer(
-                prompt,
-                padding="do_not_pad",
-                max_length=self.tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-            text_input_ids = text_inputs.input_ids.to(device)
-            self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]
-
-
-
-        # 3. prepare prompt
-        if self.prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                self.prompt,
-                device,
-                num_images_per_prompt,
-                False,
-                negative_prompt,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=None,
-                lora_scale=None,
-                clip_skip=None,
-            )
-            self.prompt_embeds = prompt_embeds
-            self.negative_prompt_embeds = negative_prompt_embeds
-
-
-
-        # 4. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
-        # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
-        # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
-        # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
-        # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
-        # operation and leads to the most reasonable results. Using the native image resolution or any other processing
-        # resolution can lead to loss of either fine details or global context in the output predictions.
-        if not skip_preprocess:
-            image, padding, original_resolution = self.image_processor.preprocess(
-                image, processing_resolution, resample_method_input, device, dtype
-            )  # [N,3,PPH,PPW]
-        else:
-            padding = (0, 0)
-            original_resolution = image.shape[2:]
-        # 5. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
-        # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
-        # Latents of each such predictions across all input images and all ensemble members are represented in the
-        # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
-        # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
-        # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
-        # code. For example, in the Marigold-LCM video processing demo, the latents initialization of a frame is taken
-        # as a convex combination of the latents output of the pipeline for the previous frame and a newly-sampled
-        # noise. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
-        # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
-        # Model invocation: self.vae.encoder.
-        image_latent, pred_latent = self.prepare_latents(
-            image, latents, generator, ensemble_size, batch_size
-        )  # [N*E,4,h,w], [N*E,4,h,w]
-
-        gaus_noise = pred_latent.detach().clone()
-        del image
-
-
-        # 6. obtain control_output
-
-        cond_scale =controlnet_conditioning_scale
-        down_block_res_samples, mid_block_res_sample = self.controlnet(
-            image_latent.detach(),
-            self.t_start,
-            encoder_hidden_states=self.prompt_embeds,
-            conditioning_scale=cond_scale,
-            guess_mode=False,
-            return_dict=False,
-        )
-
-        # 7. YOSO sampling
-        latent_x_t = self.unet(
-            pred_latent,
-            self.t_start,
-            encoder_hidden_states=self.prompt_embeds,
-            down_block_additional_residuals=down_block_res_samples,
-            mid_block_additional_residual=mid_block_res_sample,
-            return_dict=False,
-        )[0]
-
-
-        del (
-            pred_latent,
-            image_latent,
-        )
-
-        # decoder
-        prediction = self.decode_prediction(latent_x_t)
-        prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,3,PH,PW]
-
-        prediction = self.image_processor.resize_antialias(
-            prediction, original_resolution, resample_method_output, is_aa=False
-        )  # [N,3,H,W]
-        prediction = self.normalize_normals(prediction)  # [N,3,H,W]
-
-        if output_type == "np":
-            prediction = self.image_processor.pt_to_numpy(prediction)  # [N,H,W,3]
-
-        # 11. Offload all models
-        self.maybe_free_model_hooks()
-
-        return YosoNormalsOutput(
-            prediction=prediction,
-            latent=latent_x_t,
-            gaus_noise=gaus_noise,
-        )
-
-    # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents
-    def prepare_latents(
-        self,
-        image: torch.Tensor,
-        latents: Optional[torch.Tensor],
-        generator: Optional[torch.Generator],
-        ensemble_size: int,
-        batch_size: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        def retrieve_latents(encoder_output):
-            if hasattr(encoder_output, "latent_dist"):
-                return encoder_output.latent_dist.mode()
-            elif hasattr(encoder_output, "latents"):
-                return encoder_output.latents
-            else:
-                raise AttributeError("Could not access latents of provided encoder_output")
-
-
-
-        image_latent = torch.cat(
-            [
-                retrieve_latents(self.vae.encode(image[i : i + batch_size]))
-                for i in range(0, image.shape[0], batch_size)
-            ],
-            dim=0,
-        )  # [N,4,h,w]
-        image_latent = image_latent * self.vae.config.scaling_factor
-        image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]
-
-        pred_latent = torch.zeros_like(image_latent)
-        if pred_latent is None:
-            pred_latent = randn_tensor(
-                image_latent.shape,
-                generator=generator,
-                device=image_latent.device,
-                dtype=image_latent.dtype,
-            )  # [N*E,4,h,w]
-
-        return image_latent, pred_latent
-
-    def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor:
-        if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
-            raise ValueError(
-                f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
-            )
-
-        prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0]  # [B,3,H,W]
-
-        prediction = self.normalize_normals(prediction)  # [B,3,H,W]
-
-        return prediction  # [B,3,H,W]
-
-    @staticmethod
-    def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-        if normals.dim() != 4 or normals.shape[1] != 3:
-            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
-
-        norm = torch.norm(normals, dim=1, keepdim=True)
-        normals /= norm.clamp(min=eps)
-
-        return normals
-
-    @staticmethod
-    def ensemble_normals(
-        normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest"
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """
-        Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
-        the number of ensemble members for a given prediction of size `(H x W)`.
-
-        Args:
-            normals (`torch.Tensor`):
-                Input ensemble normals maps.
-            output_uncertainty (`bool`, *optional*, defaults to `False`):
-                Whether to output uncertainty map.
-            reduction (`str`, *optional*, defaults to `"closest"`):
-                Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
-                `"mean"`.
-
-        Returns:
-            A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
-            uncertainties of shape `(1, 1, H, W)`.
-        """
-        if normals.dim() != 4 or normals.shape[1] != 3:
-            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
-        if reduction not in ("closest", "mean"):
-            raise ValueError(f"Unrecognized reduction method: {reduction}.")
-
-        mean_normals = normals.mean(dim=0, keepdim=True)  # [1,3,H,W]
-        mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals)  # [1,3,H,W]
-
-        sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True)  # [E,1,H,W]
-        sim_cos = sim_cos.clamp(-1, 1)  # required to avoid NaN in uncertainty with fp16
-
-        uncertainty = None
-        if output_uncertainty:
-            uncertainty = sim_cos.arccos()  # [E,1,H,W]
-            uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi  # [1,1,H,W]
-
-        if reduction == "mean":
-            return mean_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
-
-        closest_indices = sim_cos.argmax(dim=0, keepdim=True)  # [1,1,H,W]
-        closest_indices = closest_indices.repeat(1, 3, 1, 1)  # [1,3,H,W]
-        closest_normals = torch.gather(normals, 0, closest_indices)  # [1,3,H,W]
-
-        return closest_normals, uncertainty  # [1,3,H,W], [1,1,H,W]
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps

stablenormal/scheduler/heuristics_ddimsampler.py

@@ -1,243 +0,0 @@
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput, DDIMScheduler
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.configuration_utils import register_to_config, ConfigMixin
-import pdb
-
-
-class HEURI_DDIMScheduler(DDIMScheduler, SchedulerMixin, ConfigMixin):
-
-    def set_timesteps(self, num_inference_steps: int, t_start: 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
-
-            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
-            if self.config.timestep_spacing == "linspace":
-                timesteps = (
-                    np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
-                    .round()[::-1]
-                    .copy()
-                    .astype(np.int64)
-                )
-            elif self.config.timestep_spacing == "leading":
-                step_ratio = self.config.num_train_timesteps // self.num_inference_steps
-                # creates integer timesteps by multiplying by ratio
-                # casting to int to avoid issues when num_inference_step is power of 3
-                timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
-                timesteps += self.config.steps_offset
-            elif self.config.timestep_spacing == "trailing":
-                step_ratio = self.config.num_train_timesteps / self.num_inference_steps
-                # creates integer timesteps by multiplying by ratio
-                # casting to int to avoid issues when num_inference_step is power of 3
-                timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
-                timesteps -= 1
-            else:
-                raise ValueError(
-                    f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'."
-                )
-
-            timesteps = torch.from_numpy(timesteps).to(device)
-
-
-            naive_sampling_step = num_inference_steps //2
-
-            # TODO for debug
-            # naive_sampling_step = 0
-
-            self.naive_sampling_step = naive_sampling_step
-
-            timesteps[:naive_sampling_step] = timesteps[naive_sampling_step] # refine on step 5 for 5 steps, then backward from step 6
-
-            timesteps = [timestep + 1 for timestep in timesteps]
-
-            self.timesteps = timesteps
-            self.gap = self.config.num_train_timesteps // self.num_inference_steps
-            self.prev_timesteps = [timestep for timestep in self.timesteps[1:]]
-            self.prev_timesteps.append(torch.zeros_like(self.prev_timesteps[-1]))
-
-    def step(
-            self,
-            model_output: torch.Tensor,
-            timestep: int,
-            prev_timestep: int,
-            sample: torch.Tensor,
-            eta: float = 0.0,
-            use_clipped_model_output: bool = False,
-            generator=None,
-            cur_step=None,
-            variance_noise: Optional[torch.Tensor] = None,
-            gaus_noise: Optional[torch.Tensor] = None,
-            return_dict: bool = True,
-        ) -> Union[DDIMSchedulerOutput, 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.Tensor`):
-                    The direct output from learned diffusion model.
-                timestep (`float`):
-                    The current discrete timestep in the diffusion chain.
-                pre_timestep (`float`):
-                    next_timestep
-                sample (`torch.Tensor`):
-                    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.Tensor`):
-                    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_ddim.DDIMSchedulerOutput`] or `tuple`.
-
-            Returns:
-                [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] or `tuple`:
-                    If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] 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"
-                )
-
-            # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
-            # Ideally, read DDIM paper in-detail understanding
-
-            # Notation (<variable name> -> <name in paper>
-            # - pred_noise_t -> e_theta(x_t, t)
-            # - pred_original_sample -> f_theta(x_t, t) or x_0
-            # - std_dev_t -> sigma_t
-            # - eta -> η
-            # - pred_sample_direction -> "direction pointing to x_t"
-            # - pred_prev_sample -> "x_t-1"
-
-            # 1. get previous step value (=t-1)
-
-            # trick from heuri_sampling
-            if cur_step == self.naive_sampling_step  and timestep == prev_timestep:
-                timestep += self.gap
-
-
-            prev_timestep = prev_timestep  # NOTE naive sampling
-
-            # 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
-
-            # 3. compute predicted original sample from predicted noise also called
-            # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-            if self.config.prediction_type == "epsilon":
-                pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-                pred_epsilon = model_output
-            elif self.config.prediction_type == "sample":
-                pred_original_sample = model_output
-                pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
-            elif self.config.prediction_type == "v_prediction":
-                pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
-                pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
-            else:
-                raise ValueError(
-                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
-                    " `v_prediction`"
-                )
-
-            # 4. Clip or threshold "predicted x_0"
-            if self.config.thresholding:
-                pred_original_sample = self._threshold_sample(pred_original_sample)
-
-            # 5. compute variance: "sigma_t(η)" -> see formula (16)
-            # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
-            variance = self._get_variance(timestep, prev_timestep)
-            std_dev_t = eta * variance ** (0.5)
-
-
-            if use_clipped_model_output:
-                # the pred_epsilon is always re-derived from the clipped x_0 in Glide
-                pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
-
-            # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-            pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon
-
-            # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-            prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
-
-            if eta > 0:
-                if variance_noise is not None and generator is not None:
-                    raise ValueError(
-                        "Cannot pass both generator and variance_noise. Please make sure that either `generator` or"
-                        " `variance_noise` stays `None`."
-                    )
-
-                if variance_noise is None:
-                    variance_noise = randn_tensor(
-                        model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype
-                    )
-                variance = std_dev_t * variance_noise
-
-                prev_sample = prev_sample + variance
-
-            if cur_step < self.naive_sampling_step:
-                prev_sample = self.add_noise(pred_original_sample, torch.randn_like(pred_original_sample), timestep)
-
-            if not return_dict:
-                return (prev_sample,)
-
-
-            return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
-
-
-
-    def add_noise(
-        self,
-        original_samples: torch.Tensor,
-        noise: torch.Tensor,
-        timesteps: torch.IntTensor,
-    ) -> torch.Tensor:
-        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
-        # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
-        # for the subsequent add_noise calls
-        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
-        alphas_cumprod = self.alphas_cumprod.to(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

stablenormal/stablecontrolnet.py

@@ -1,1354 +0,0 @@
-# Copyright 2024 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.
-
-
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import PIL.Image
-import torch
-import torch.nn.functional as F
-from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
-
-from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...image_processor import PipelineImageInput, VaeImageProcessor
-from ...loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin
-from ...models import AutoencoderKL, ControlNetModel, ImageProjection, UNet2DConditionModel
-from ...models.lora import adjust_lora_scale_text_encoder
-from ...schedulers import KarrasDiffusionSchedulers
-from ...utils import (
-    USE_PEFT_BACKEND,
-    deprecate,
-    logging,
-    replace_example_docstring,
-    scale_lora_layers,
-    unscale_lora_layers,
-)
-from ...utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor
-from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin
-from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
-from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-from .multicontrolnet import MultiControlNetModel
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> # !pip install opencv-python transformers accelerate
-        >>> from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
-        >>> from diffusers.utils import load_image
-        >>> import numpy as np
-        >>> import torch
-
-        >>> import cv2
-        >>> from PIL import Image
-
-        >>> # download an image
-        >>> image = load_image(
-        ...     "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
-        ... )
-        >>> image = np.array(image)
-
-        >>> # get canny image
-        >>> image = cv2.Canny(image, 100, 200)
-        >>> image = image[:, :, None]
-        >>> image = np.concatenate([image, image, image], axis=2)
-        >>> canny_image = Image.fromarray(image)
-
-        >>> # load control net and stable diffusion v1-5
-        >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-        >>> pipe = StableDiffusionControlNetPipeline.from_pretrained(
-        ...     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
-        ... )
-
-        >>> # speed up diffusion process with faster scheduler and memory optimization
-        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-        >>> # remove following line if xformers is not installed
-        >>> pipe.enable_xformers_memory_efficient_attention()
-
-        >>> pipe.enable_model_cpu_offload()
-
-        >>> # generate image
-        >>> generator = torch.manual_seed(0)
-        >>> image = pipe(
-        ...     "futuristic-looking woman", num_inference_steps=20, generator=generator, image=canny_image
-        ... ).images[0]
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class StableDiffusionControlNetPipeline(
-    DiffusionPipeline,
-    StableDiffusionMixin,
-    TextualInversionLoaderMixin,
-    LoraLoaderMixin,
-    IPAdapterMixin,
-    FromSingleFileMixin,
-):
-    r"""
-    Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    The pipeline also inherits the following loading methods:
-        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
-        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
-        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
-        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
-        - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
-
-    Args:
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
-        text_encoder ([`~transformers.CLIPTextModel`]):
-            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
-        tokenizer ([`~transformers.CLIPTokenizer`]):
-            A `CLIPTokenizer` to tokenize text.
-        unet ([`UNet2DConditionModel`]):
-            A `UNet2DConditionModel` to denoise the encoded image latents.
-        controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
-            Provides additional conditioning to the `unet` during the denoising process. If you set multiple
-            ControlNets as a list, the outputs from each ControlNet are added together to create one combined
-            additional conditioning.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
-            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
-        safety_checker ([`StableDiffusionSafetyChecker`]):
-            Classification module that estimates whether generated images could be considered offensive or harmful.
-            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
-            about a model's potential harms.
-        feature_extractor ([`~transformers.CLIPImageProcessor`]):
-            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
-    """
-
-    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
-    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
-    _exclude_from_cpu_offload = ["safety_checker"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        tokenizer: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
-        scheduler: KarrasDiffusionSchedulers,
-        safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
-        image_encoder: CLIPVisionModelWithProjection = None,
-        requires_safety_checker: bool = True,
-    ):
-        super().__init__()
-
-        if safety_checker is None and requires_safety_checker:
-            logger.warning(
-                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
-                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
-                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
-                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
-                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
-                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
-            )
-
-        if safety_checker is not None and feature_extractor is None:
-            raise ValueError(
-                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
-                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
-            )
-
-        if isinstance(controlnet, (list, tuple)):
-            controlnet = MultiControlNetModel(controlnet)
-
-        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,
-            image_encoder=image_encoder,
-        )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)
-        self.control_image_processor = VaeImageProcessor(
-            vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False
-        )
-        self.register_to_config(requires_safety_checker=requires_safety_checker)
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
-    def _encode_prompt(
-        self,
-        prompt,
-        device,
-        num_images_per_prompt,
-        do_classifier_free_guidance,
-        negative_prompt=None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        lora_scale: Optional[float] = None,
-        **kwargs,
-    ):
-        deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
-        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
-
-        prompt_embeds_tuple = self.encode_prompt(
-            prompt=prompt,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            lora_scale=lora_scale,
-            **kwargs,
-        )
-
-        # concatenate for backwards comp
-        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
-
-        return prompt_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt,
-        device,
-        num_images_per_prompt,
-        do_classifier_free_guidance,
-        negative_prompt=None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        lora_scale: Optional[float] = None,
-        clip_skip: Optional[int] = 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
-            do_classifier_free_guidance (`bool`):
-                whether to use classifier free guidance or not
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            prompt_embeds (`torch.Tensor`, *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.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            lora_scale (`float`, *optional*):
-                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-        """
-        # set lora scale so that monkey patched LoRA
-        # function of text encoder can correctly access it
-        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
-            self._lora_scale = lora_scale
-
-            # dynamically adjust the LoRA scale
-            if not USE_PEFT_BACKEND:
-                adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
-            else:
-                scale_lora_layers(self.text_encoder, lora_scale)
-
-        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]
-
-        if prompt_embeds is None:
-            # textual inversion: process multi-vector tokens if necessary
-            if isinstance(self, TextualInversionLoaderMixin):
-                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
-
-            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
-
-            if clip_skip is None:
-                prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
-                prompt_embeds = prompt_embeds[0]
-            else:
-                prompt_embeds = self.text_encoder(
-                    text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
-                )
-                # Access the `hidden_states` first, that contains a tuple of
-                # all the hidden states from the encoder layers. Then index into
-                # the tuple to access the hidden states from the desired layer.
-                prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
-                # We also need to apply the final LayerNorm here to not mess with the
-                # representations. The `last_hidden_states` that we typically use for
-                # obtaining the final prompt representations passes through the LayerNorm
-                # layer.
-                prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)
-
-        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)
-
-        # get unconditional embeddings for classifier free guidance
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            uncond_tokens: List[str]
-            if negative_prompt is None:
-                uncond_tokens = [""] * batch_size
-            elif prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif isinstance(negative_prompt, str):
-                uncond_tokens = [negative_prompt]
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-            else:
-                uncond_tokens = negative_prompt
-
-            # textual inversion: process multi-vector tokens if necessary
-            if isinstance(self, TextualInversionLoaderMixin):
-                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
-
-            max_length = prompt_embeds.shape[1]
-            uncond_input = self.tokenizer(
-                uncond_tokens,
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-
-            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
-                attention_mask = uncond_input.attention_mask.to(device)
-            else:
-                attention_mask = None
-
-            negative_prompt_embeds = self.text_encoder(
-                uncond_input.input_ids.to(device),
-                attention_mask=attention_mask,
-            )
-            negative_prompt_embeds = negative_prompt_embeds[0]
-
-        if do_classifier_free_guidance:
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-
-            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        if self.text_encoder is not None:
-            if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
-                # Retrieve the original scale by scaling back the LoRA layers
-                unscale_lora_layers(self.text_encoder, lora_scale)
-
-        return prompt_embeds, negative_prompt_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
-    def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
-        dtype = next(self.image_encoder.parameters()).dtype
-
-        if not isinstance(image, torch.Tensor):
-            image = self.feature_extractor(image, return_tensors="pt").pixel_values
-
-        image = image.to(device=device, dtype=dtype)
-        if output_hidden_states:
-            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
-            image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
-            uncond_image_enc_hidden_states = self.image_encoder(
-                torch.zeros_like(image), output_hidden_states=True
-            ).hidden_states[-2]
-            uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
-                num_images_per_prompt, dim=0
-            )
-            return image_enc_hidden_states, uncond_image_enc_hidden_states
-        else:
-            image_embeds = self.image_encoder(image).image_embeds
-            image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-            uncond_image_embeds = torch.zeros_like(image_embeds)
-
-            return image_embeds, uncond_image_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
-    def prepare_ip_adapter_image_embeds(
-        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
-    ):
-        if ip_adapter_image_embeds is None:
-            if not isinstance(ip_adapter_image, list):
-                ip_adapter_image = [ip_adapter_image]
-
-            if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
-                raise ValueError(
-                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
-                )
-
-            image_embeds = []
-            for single_ip_adapter_image, image_proj_layer in zip(
-                ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
-            ):
-                output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
-                single_image_embeds, single_negative_image_embeds = self.encode_image(
-                    single_ip_adapter_image, device, 1, output_hidden_state
-                )
-                single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
-                single_negative_image_embeds = torch.stack(
-                    [single_negative_image_embeds] * num_images_per_prompt, dim=0
-                )
-
-                if do_classifier_free_guidance:
-                    single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
-                    single_image_embeds = single_image_embeds.to(device)
-
-                image_embeds.append(single_image_embeds)
-        else:
-            repeat_dims = [1]
-            image_embeds = []
-            for single_image_embeds in ip_adapter_image_embeds:
-                if do_classifier_free_guidance:
-                    single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
-                    single_image_embeds = single_image_embeds.repeat(
-                        num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
-                    )
-                    single_negative_image_embeds = single_negative_image_embeds.repeat(
-                        num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
-                    )
-                    single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
-                else:
-                    single_image_embeds = single_image_embeds.repeat(
-                        num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
-                    )
-                image_embeds.append(single_image_embeds)
-
-        return image_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
-    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
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
-    def decode_latents(self, latents):
-        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
-        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
-
-        latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents, return_dict=False)[0]
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        return image
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def check_inputs(
-        self,
-        prompt,
-        image,
-        callback_steps,
-        negative_prompt=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        ip_adapter_image=None,
-        ip_adapter_image_embeds=None,
-        controlnet_conditioning_scale=1.0,
-        control_guidance_start=0.0,
-        control_guidance_end=1.0,
-        callback_on_step_end_tensor_inputs=None,
-    ):
-        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
-            raise ValueError(
-                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
-                f" {type(callback_steps)}."
-            )
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}."
-                )
-
-        # Check `image`
-        is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
-            self.controlnet, torch._dynamo.eval_frame.OptimizedModule
-        )
-        if (
-            isinstance(self.controlnet, ControlNetModel)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, ControlNetModel)
-        ):
-            self.check_image(image, prompt, prompt_embeds)
-        elif (
-            isinstance(self.controlnet, MultiControlNetModel)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
-        ):
-            if not isinstance(image, list):
-                raise TypeError("For multiple controlnets: `image` must be type `list`")
-
-            # When `image` is a nested list:
-            # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])
-            elif any(isinstance(i, list) for i in image):
-                transposed_image = [list(t) for t in zip(*image)]
-                if len(transposed_image) != len(self.controlnet.nets):
-                    raise ValueError(
-                        f"For multiple controlnets: if you pass`image` as a list of list, each sublist must have the same length as the number of controlnets, but the sublists in `image` got {len(transposed_image)} images and {len(self.controlnet.nets)} ControlNets."
-                    )
-                for image_ in transposed_image:
-                    self.check_image(image_, prompt, prompt_embeds)
-            elif len(image) != len(self.controlnet.nets):
-                raise ValueError(
-                    f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {len(self.controlnet.nets)} ControlNets."
-                )
-            else:
-                for image_ in image:
-                    self.check_image(image_, prompt, prompt_embeds)
-        else:
-            assert False
-
-        # Check `controlnet_conditioning_scale`
-        if (
-            isinstance(self.controlnet, ControlNetModel)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, ControlNetModel)
-        ):
-            if not isinstance(controlnet_conditioning_scale, float):
-                raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
-        elif (
-            isinstance(self.controlnet, MultiControlNetModel)
-            or is_compiled
-            and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
-        ):
-            if isinstance(controlnet_conditioning_scale, list):
-                if any(isinstance(i, list) for i in controlnet_conditioning_scale):
-                    raise ValueError(
-                        "A single batch of varying conditioning scale settings (e.g. [[1.0, 0.5], [0.2, 0.8]]) is not supported at the moment. "
-                        "The conditioning scale must be fixed across the batch."
-                    )
-            elif isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(
-                self.controlnet.nets
-            ):
-                raise ValueError(
-                    "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"
-                    " the same length as the number of controlnets"
-                )
-        else:
-            assert False
-
-        if not isinstance(control_guidance_start, (tuple, list)):
-            control_guidance_start = [control_guidance_start]
-
-        if not isinstance(control_guidance_end, (tuple, list)):
-            control_guidance_end = [control_guidance_end]
-
-        if len(control_guidance_start) != len(control_guidance_end):
-            raise ValueError(
-                f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list."
-            )
-
-        if isinstance(self.controlnet, MultiControlNetModel):
-            if len(control_guidance_start) != len(self.controlnet.nets):
-                raise ValueError(
-                    f"`control_guidance_start`: {control_guidance_start} has {len(control_guidance_start)} elements but there are {len(self.controlnet.nets)} controlnets available. Make sure to provide {len(self.controlnet.nets)}."
-                )
-
-        for start, end in zip(control_guidance_start, control_guidance_end):
-            if start >= end:
-                raise ValueError(
-                    f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
-                )
-            if start < 0.0:
-                raise ValueError(f"control guidance start: {start} can't be smaller than 0.")
-            if end > 1.0:
-                raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")
-
-        if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
-            raise ValueError(
-                "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
-            )
-
-        if ip_adapter_image_embeds is not None:
-            if not isinstance(ip_adapter_image_embeds, list):
-                raise ValueError(
-                    f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
-                )
-            elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
-                raise ValueError(
-                    f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
-                )
-
-    def check_image(self, image, prompt, prompt_embeds):
-        image_is_pil = isinstance(image, PIL.Image.Image)
-        image_is_tensor = isinstance(image, torch.Tensor)
-        image_is_np = isinstance(image, np.ndarray)
-        image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)
-        image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)
-        image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)
-
-        if (
-            not image_is_pil
-            and not image_is_tensor
-            and not image_is_np
-            and not image_is_pil_list
-            and not image_is_tensor_list
-            and not image_is_np_list
-        ):
-            raise TypeError(
-                f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}"
-            )
-
-        if image_is_pil:
-            image_batch_size = 1
-        else:
-            image_batch_size = len(image)
-
-        if prompt is not None and isinstance(prompt, str):
-            prompt_batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            prompt_batch_size = len(prompt)
-        elif prompt_embeds is not None:
-            prompt_batch_size = prompt_embeds.shape[0]
-
-        if image_batch_size != 1 and image_batch_size != prompt_batch_size:
-            raise ValueError(
-                f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"
-            )
-
-    def prepare_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=torch.float32)
-        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
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
-    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-        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."
-            )
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        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
-
-    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
-    def get_guidance_scale_embedding(
-        self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
-    ) -> torch.Tensor:
-        """
-        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-
-        Args:
-            w (`torch.Tensor`):
-                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
-            embedding_dim (`int`, *optional*, defaults to 512):
-                Dimension of the embeddings to generate.
-            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
-                Data type of the generated embeddings.
-
-        Returns:
-            `torch.Tensor`: Embedding vectors with shape `(len(w), 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
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def clip_skip(self):
-        return self._clip_skip
-
-    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-    # corresponds to doing no classifier free guidance.
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
-
-    @property
-    def cross_attention_kwargs(self):
-        return self._cross_attention_kwargs
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        image: PipelineImageInput = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        timesteps: List[int] = None,
-        sigmas: List[float] = None,
-        guidance_scale: float = 7.5,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        ip_adapter_image: Optional[PipelineImageInput] = None,
-        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = 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]] = 1.0,
-        guess_mode: bool = False,
-        control_guidance_start: Union[float, List[float]] = 0.0,
-        control_guidance_end: Union[float, List[float]] = 1.0,
-        clip_skip: Optional[int] = None,
-        callback_on_step_end: Optional[
-            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
-        ] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        **kwargs,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
-            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
-                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
-                The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
-                specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
-                as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
-                width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
-                images must be passed as a list such that each element of the list can be correctly batched for input
-                to a single ControlNet. When `prompt` is a list, and if a list of images is passed for a single
-                ControlNet, each will be paired with each prompt in the `prompt` list. This also applies to multiple
-                ControlNets, where a list of image lists can be passed to batch for each prompt and each ControlNet.
-            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
-                The width in pixels of the generated image.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
-                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
-                passed will be used. Must be in descending order.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to 7.5):
-                A higher guidance scale value encourages the model to generate images closely linked to the text
-                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
-                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            eta (`float`, *optional*, defaults to 0.0):
-                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
-                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
-                provided, text embeddings are generated from the `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
-                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
-            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
-            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
-                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
-                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
-                contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
-                provided, embeddings are computed from the `ip_adapter_image` input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
-                plain tuple.
-            callback (`Callable`, *optional*):
-                A function that calls every `callback_steps` steps during inference. The function is called with the
-                following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
-            callback_steps (`int`, *optional*, defaults to 1):
-                The frequency at which the `callback` function is called. If not specified, the callback is called at
-                every step.
-            cross_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
-                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
-                The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
-                to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
-                the corresponding scale as a list.
-            guess_mode (`bool`, *optional*, defaults to `False`):
-                The ControlNet encoder tries to recognize the content of the input image even if you remove all
-                prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
-            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
-                The percentage of total steps at which the ControlNet starts applying.
-            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
-                The percentage of total steps at which the ControlNet stops applying.
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
-                otherwise a `tuple` is returned where the first element is a list with the generated images and the
-                second element is a list of `bool`s indicating whether the corresponding generated image contains
-                "not-safe-for-work" (nsfw) content.
-        """
-
-        callback = kwargs.pop("callback", None)
-        callback_steps = kwargs.pop("callback_steps", None)
-
-        if callback is not None:
-            deprecate(
-                "callback",
-                "1.0.0",
-                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
-            )
-        if callback_steps is not None:
-            deprecate(
-                "callback_steps",
-                "1.0.0",
-                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
-            )
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
-
-        # align format for control guidance
-        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],
-            )
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            image,
-            callback_steps,
-            negative_prompt,
-            prompt_embeds,
-            negative_prompt_embeds,
-            ip_adapter_image,
-            ip_adapter_image_embeds,
-            controlnet_conditioning_scale,
-            control_guidance_start,
-            control_guidance_end,
-            callback_on_step_end_tensor_inputs,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._clip_skip = clip_skip
-        self._cross_attention_kwargs = cross_attention_kwargs
-
-        # 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
-
-        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):
-            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)
-
-        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
-        text_encoder_lora_scale = (
-            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
-        )
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            prompt,
-            device,
-            num_images_per_prompt,
-            self.do_classifier_free_guidance,
-            negative_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            lora_scale=text_encoder_lora_scale,
-            clip_skip=self.clip_skip,
-        )
-        # For classifier free guidance, we need to do two forward passes.
-        # Here we concatenate the unconditional and text embeddings into a single batch
-        # to avoid doing two forward passes
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-
-        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
-            image_embeds = self.prepare_ip_adapter_image_embeds(
-                ip_adapter_image,
-                ip_adapter_image_embeds,
-                device,
-                batch_size * num_images_per_prompt,
-                self.do_classifier_free_guidance,
-            )
-
-        # 4. Prepare image
-        if isinstance(controlnet, ControlNetModel):
-            image = self.prepare_image(
-                image=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=self.do_classifier_free_guidance,
-                guess_mode=guess_mode,
-            )
-            height, width = image.shape[-2:]
-        elif isinstance(controlnet, MultiControlNetModel):
-            images = []
-
-            # Nested lists as ControlNet condition
-            if isinstance(image[0], list):
-                # Transpose the nested image list
-                image = [list(t) for t in zip(*image)]
-
-            for image_ in image:
-                image_ = self.prepare_image(
-                    image=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=self.do_classifier_free_guidance,
-                    guess_mode=guess_mode,
-                )
-
-                images.append(image_)
-
-            image = images
-            height, width = image[0].shape[-2:]
-        else:
-            assert False
-
-        # 5. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device, timesteps, sigmas
-        )
-        self._num_timesteps = len(timesteps)
-
-        # 6. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-
-        # 6.5 Optionally get Guidance Scale Embedding
-        timestep_cond = None
-        if self.unet.config.time_cond_proj_dim is not None:
-            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
-            timestep_cond = self.get_guidance_scale_embedding(
-                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
-            ).to(device=device, dtype=latents.dtype)
-
-        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7.1 Add image embeds for IP-Adapter
-        added_cond_kwargs = (
-            {"image_embeds": image_embeds}
-            if ip_adapter_image is not None or ip_adapter_image_embeds is not None
-            else None
-        )
-
-        # 7.2 Create tensor stating which controlnets to keep
-        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)
-
-        # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        is_unet_compiled = is_compiled_module(self.unet)
-        is_controlnet_compiled = is_compiled_module(self.controlnet)
-        is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # Relevant thread:
-                # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428
-                if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1:
-                    torch._inductor.cudagraph_mark_step_begin()
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                # controlnet(s) inference
-                if guess_mode and self.do_classifier_free_guidance:
-                    # Infer ControlNet only for the conditional batch.
-                    control_model_input = latents
-                    control_model_input = self.scheduler.scale_model_input(control_model_input, t)
-                    controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
-                else:
-                    control_model_input = latent_model_input
-                    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,
-                    t,
-                    encoder_hidden_states=controlnet_prompt_embeds,
-                    controlnet_cond=image,
-                    conditioning_scale=cond_scale,
-                    guess_mode=guess_mode,
-                    return_dict=False,
-                )
-
-                if guess_mode and self.do_classifier_free_guidance:
-                    # Infered ControlNet only for the conditional batch.
-                    # To apply the output of ControlNet to both the unconditional and conditional batches,
-                    # add 0 to the unconditional batch to keep it unchanged.
-                    down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
-                    mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
-
-                # predict the noise residual
-                noise_pred = self.unet(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep_cond=timestep_cond,
-                    cross_attention_kwargs=self.cross_attention_kwargs,
-                    down_block_additional_residuals=down_block_res_samples,
-                    mid_block_additional_residual=mid_block_res_sample,
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # perform guidance
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        step_idx = i // getattr(self.scheduler, "order", 1)
-                        callback(step_idx, t, latents)
-
-        # If we do sequential model offloading, let's offload unet and controlnet
-        # manually for max memory savings
-        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(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
-                0
-            ]
-            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
-        else:
-            image = latents
-            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)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image, has_nsfw_concept)
-
-        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)