chore: import upstream snapshot with attribution

2026-07-13 13:05:14 +08:00
commit 2a547be7fe
7904 changed files with 1000926 additions and 0 deletions
@@ -0,0 +1,92 @@
+<!--[metadata]
+title = "ControlNet"
+tags = ["ControlNet", "Canny", "Hugging Face", "Stable diffusion", "Tensor", "Text"]
+thumbnail = "https://static.rerun.io/controlnet/2e984b27dd8120fb89d4e805df9da506ea6d9138/480w.png"
+thumbnail_dimensions = [480, 480]
+-->
+
+Use [Hugging Face's ControlNet](https://huggingface.co/docs/diffusers/using-diffusers/controlnet#controlnet) to generate an image from text, conditioned on detected edges from another image.
+
+<picture>
+  <source media="(max-width: 480px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/480w.png">
+  <source media="(max-width: 768px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/768w.png">
+  <source media="(max-width: 1024px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/1024w.png">
+  <source media="(max-width: 1200px)" srcset="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/1200w.png">
+  <img src="https://static.rerun.io/controlnet/8aace9c59a423c2eeabe4b7f9abb5187559c52e8/full.png" alt="">
+</picture>
+
+
+## Used Rerun types
+[`Image`](https://www.rerun.io/docs/reference/types/archetypes/image), [`Tensor`](https://www.rerun.io/docs/reference/types/archetypes/tensor), [`TextDocument`](https://www.rerun.io/docs/reference/types/archetypes/text_document)
+
+
+## Background
+[Hugging Face's ControlNet](https://huggingface.co/docs/diffusers/using-diffusers/controlnet#controlnet) allows to condition Stable Diffusion on various modalities. In this example we condition on edges detected by the Canny edge detector to keep our shape intact while generating an image with Stable Diffusion. This generates a new image based on the text prompt, but that still retains the structure of the control image. We visualize the whole generation process with Rerun.
+
+https://vimeo.com/870289439?autoplay=1&loop=1&autopause=0&background=1&muted=1&ratio=1440:1080
+
+## Logging and visualizing with Rerun
+The visualizations in this example were created with the following Rerun code.
+
+### Images
+```python
+rr.log("input/raw", rr.Image(image), static=True)
+rr.log("input/canny", rr.Image(canny_image), static=True)
+```
+The input image and control canny_image are marked as static and logged in rerun.
+
+Static entities belong to all timelines (existing ones, and ones not yet created) and are shown leftmost in the time panel in the viewer. This is useful for entities that aren't part of normal data capture, but set the scene for how they are shown.
+
+This designation ensures their constant availability across all timelines in Rerun, aiding in consistent comparison and documentation.
+
+### Prompts
+```python
+rr.log("positive_prompt", rr.TextDocument(prompt), static=True)
+rr.log("negative_prompt", rr.TextDocument(negative_prompt), static=True)
+```
+The positive and negative prompt used for generation is logged to Rerun.
+
+### Custom diffusion step callback
+We use a custom callback function for ControlNet that logs the output and the latent values at each timestep, which makes it possible for us to view all timesteps of the generation in Rerun.
+```python
+def controlnet_callback(
+    iteration: int, timestep: float, latents: torch.Tensor, pipeline: StableDiffusionXLControlNetPipeline
+) -> None:
+    rr.set_time("iteration", sequence=iteration)
+    rr.set_time("timestep", duration=timestep)
+    rr.log("output", rr.Image(image))
+    rr.log("latent", rr.Tensor(latents.squeeze(), dim_names=["channel", "height", "width"]))
+```
+
+### Output image
+```python
+rr.log("output", rr.Image(images))
+```
+Finally we log the output image generated by ControlNet.
+
+## Run the code
+
+To run this example, make sure you have the Rerun repository checked out and the latest SDK installed:
+```bash
+pip install --upgrade rerun-sdk  # install the latest Rerun SDK
+git clone git@github.com:rerun-io/rerun.git  # Clone the repository
+cd rerun
+git checkout latest  # Check out the commit matching the latest SDK release
+```
+
+Install the necessary libraries specified in the requirements file:
+```bash
+pip install -e examples/python/controlnet
+```
+
+To experiment with the provided example, simply execute the main Python script:
+```bash
+python -m controlnet
+```
+
+You can specify your own image and prompts using
+```bash
+python -m controlnet [--img-path IMG_PATH] [--prompt PROMPT] [--negative-prompt NEGATIVE_PROMPT]
+```
+
+This example requires a machine with CUDA backend available for pytorch (GPU) to work.
@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+"""
+Example running ControlNet conditioned on Canny edges.
+
+Based on <https://huggingface.co/docs/diffusers/using-diffusers/controlnet>.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Any
+
+import cv2
+import numpy as np
+import PIL.Image
+import requests
+import torch
+from diffusers import (
+    AutoencoderKL,
+    ControlNetModel,
+    StableDiffusionXLControlNetPipeline,
+)
+
+import rerun as rr
+import rerun.blueprint as rrb
+
+RERUN_LOGO_URL = "https://storage.googleapis.com/rerun-example-datasets/controlnet/rerun-icon-1000.png"
+
+
+def controlnet_callback(
+    pipe: StableDiffusionXLControlNetPipeline,
+    step_index: int,
+    timestep: float,
+    callback_kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    rr.set_time("iteration", sequence=step_index)
+    rr.set_time("timestep", duration=timestep)
+    latents = callback_kwargs["latents"]
+
+    image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]  # type: ignore[attr-defined]
+    image = pipe.image_processor.postprocess(image, output_type="np").squeeze()  # type: ignore[attr-defined]
+    rr.log("output", rr.Image(image))
+    rr.log("latent", rr.Tensor(latents.squeeze(), dim_names=["channel", "height", "width"]))
+
+    return callback_kwargs
+
+
+def run_canny_controlnet(image_path: str, prompt: str, negative_prompt: str) -> None:
+    if not torch.cuda.is_available():
+        print("This example requires a torch with CUDA, but no CUDA device found. Aborting.")
+        return
+
+    if image_path.startswith(("http://", "https://")):
+        pil_image = PIL.Image.open(requests.get(image_path, stream=True).content)
+    elif os.path.isfile(image_path):
+        pil_image = PIL.Image.open(image_path)
+    else:
+        raise ValueError(f"Invalid image_path: {image_path}")
+
+    image = np.array(pil_image)
+
+    if image.shape[2] == 4:  # RGBA image
+        rgb_image = image[..., :3]  # RGBA to RGB
+        rgb_image[image[..., 3] < 200] = 0.0  # reduces artifacts for transparent parts
+    else:
+        rgb_image = image
+
+    low_threshold = 100.0
+    high_threshold = 200.0
+    canny_data = cv2.Canny(rgb_image, low_threshold, high_threshold)
+    canny_data = canny_data[:, :, None]
+    # cv2.dilate(kjgk
+    canny_data = np.concatenate([canny_data, canny_data, canny_data], axis=2)
+    canny_image = PIL.Image.fromarray(canny_data)
+
+    rr.log("input/raw", rr.Image(image), static=True)
+    rr.log("input/canny", rr.Image(canny_image), static=True)
+
+    controlnet = ControlNetModel.from_pretrained(
+        "diffusers/controlnet-canny-sdxl-1.0",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        controlnet=controlnet,
+        vae=vae,
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+
+    pipeline.enable_model_cpu_offload()
+
+    rr.log("positive_prompt", rr.TextDocument(prompt), static=True)
+    rr.log("negative_prompt", rr.TextDocument(negative_prompt), static=True)
+
+    images = pipeline(
+        prompt,
+        negative_prompt=negative_prompt,
+        image=canny_image,  # add batch dimension
+        controlnet_conditioning_scale=0.5,
+        callback_on_step_end=controlnet_callback,
+    ).images[0]
+
+    rr.log("output", rr.Image(images))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Use Canny-conditioned ControlNet to generate image.")
+    parser.add_argument(
+        "--img-path",
+        type=str,
+        help="Path to image used as input for Canny edge detector.",
+        default=RERUN_LOGO_URL,
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        help="Prompt used as input for ControlNet.",
+        default="aerial view, a futuristic research complex in a bright foggy jungle, hard lighting",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        help="Negative prompt used as input for ControlNet.",
+        default="low quality, bad quality, sketches",
+    )
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+
+    rr.script_setup(
+        args,
+        "rerun_example_controlnet",
+        default_blueprint=rrb.Horizontal(
+            rrb.Grid(
+                rrb.Spatial2DView(origin="input/raw"),
+                rrb.Spatial2DView(origin="input/canny"),
+                rrb.Vertical(
+                    rrb.TextDocumentView(origin="positive_prompt"),
+                    rrb.TextDocumentView(origin="negative_prompt"),
+                ),
+                rrb.TensorView(origin="latent"),
+            ),
+            rrb.Spatial2DView(origin="output"),
+        ),
+    )
+    run_canny_controlnet(args.img_path, args.prompt, args.negative_prompt)
+    rr.script_teardown(args)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,22 @@
+[project]
+name = "controlnet"
+version = "0.1.0"
+requires-python = ">=3.10"
+readme = "README.md"
+dependencies = [
+  "accelerate",
+  "opencv-python",
+  "pillow",
+  "diffusers<0.39",
+  "numpy",
+  "torch", # this will use the version defined in the uv workspace
+  "transformers",
+  "rerun-sdk",
+]
+
+[project.scripts]
+controlnet = "controlnet:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"