SWAN: Generative AI for Safe and Photorealistic Drone Light Shows

Pascal Reinhold, Alexander Gräfe, and Sebastian Trimpe

SWAN is a generative AI pipeline that creates dynamic choreographies for drone light shows with thousands of UAVs, based on natural language descriptions.

🎥 Video Demo

Watch the pipeline in action:

---

💻 System Requirements

To run SWAN effectively, your system should meet the following specifications:

• GPU: NVIDIA GPU with at least 16 GB VRAM
• RAM: 32 GB minimum recommended
• Storage: 100 GB of free disk space

🛠️ Setup & Installation

1. Clone the repository

   Make sure to use the --recursive flag to include the ComfyUI submodule:

   git clone --recursive git@github.com:Data-Science-in-Mechanical-Engineering/SWAN.git
   cd SWAN

   Note: Models will be cached in ./weights and ./comfyui_models on your host filesystem.

2. Build the Docker image

   docker build -t swan .

3. Download Model Weights

   You will be prompted to choose 1 (for core models only) or 2 (for all models, including video models).

   docker compose run swan python download_weights.py

🚀 Usage

To start the pipeline and launch the Gradio user interface:

docker compose run swan python main.py

Once your animation is generated, you can find the output files in the ./out/<name>/ directory.

🧠 Codebase Overview

The SWAN pipeline is orchestrated from main.py (CLI/local UI entry point) and implemented as a set of modular stages under the swan/ package.

🏗️ Architecture & Environment

Component	Description
Docker Environment	Based on nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 with a UV-managed Python 3.12 virtual environment (/opt/venv). docker-compose.yaml handles volume mounts (/app, ./weights, ./comfyui_models) and GPU passthrough.
Models & Assets	Weights are fetched via download_weights.py. ComfyUI is included via a Git submodule. Static configs (prompts, workflow JSONs, AXSwarm settings) live in static_data/.
Entrance Point	main.py starts the app.
SWAN Pipeline	The main pipeline is implemented in (swan/pipeline.py).
User Interfaces	The UI is implemented in ui/ui.py.
Utilities	swan/utils.py contains shared visualization helpers and rendering tools.

⚙️ Pipeline Stages

1. Video Generation (swan/video_generation.py)

• Prompting: Expands user prompts via a local LLM (PromptExpander) into cinematic components and a segmentation prompt.
• Generation: Produces a start frame and follow-up video using ComfyUIServer / ComfyUIClient wrappers around the ComfyUI HTTP API. Workflow templates are executed asynchronously.

2. Segmentation (swan/tracking.py)

• Segments the main object frame-by-frame using LangSAM + SAM 2.1.
• Results are cached for efficiency. Runs separately in CLI mode, or integrated into the tracking stage in UI mode.

3. Tracking (swan/tracking.py)

• Initialization: Samples initial drone positions via Centroidal Voronoi Tessellation (K-Means) inside the segmentation mask.
• Motion: Tracks points through the video with CoTracker3.
• Refinement: Applies a force field to maintain mask boundaries, prevent overcrowding, heal lost tracks, and remove outliers. Outputs 2D trajectories and visibility flags.

4. Trajectory Generation (swan/trajectory_generation.py)

• Translation: Smooths trajectories with splines and scales image-space units to real-world meters, ensuring inter-drone distances meet safety margins.
• Optimization: Solves a min-cost-flow problem (networkx) to assign segments to the minimum number of physical drones, injecting extra drones where kinematically necessary.
• Routing: Computes 3D takeoff/landing pads, routes collision-free transitions (via X-axis avoidance bumps), and finalizes quintic B-splines. Outputs to initial_trajectories.npz.

5. Simulation & Safety Filter (swan/simulation.py)

• Solver: Runs the AXSwarm MPC solver inside a crazyflow simulation.
• Validation: Uses a custom CollisionlessSim to avoid contact physics and save memory. Verifies dense trajectories for collisions and outputs to simulation_results.npz.

📖 Citation

If you find this codebase or our research useful, please consider citing our paper:

@misc{reinhold2026generativeaisafephotorealistic,
title={Generative AI for Safe and Photorealistic Drone Light Shows},
author={Pascal Reinhold and Alexander Gräfe and Sebastian Trimpe},
year={2026},
url={https://arxiv.org/abs/2606.25458},
}