File size: 4,991 Bytes
0261f7a
40440d6
 
 
 
 
 
 
 
0261f7a
40440d6
 
 
 
 
 
 
 
0261f7a
 
 
 
 
40440d6
 
 
 
 
 
0261f7a
40440d6
 
 
 
 
0261f7a
40440d6
 
0261f7a
40440d6
0261f7a
40440d6
 
 
 
 
 
 
 
 
 
0261f7a
40440d6
 
 
 
 
 
 
 
 
0261f7a
40440d6
0261f7a
40440d6
0261f7a
 
 
40440d6
 
 
 
0261f7a
40440d6
 
0261f7a
40440d6
0261f7a
40440d6
0261f7a
 
 
 
 
 
 
 
 
 
40440d6
0261f7a
40440d6
0261f7a
40440d6
0261f7a
40440d6
0261f7a
40440d6
 
 
 
 
 
 
0261f7a
40440d6
 
0261f7a
40440d6
0261f7a
40440d6
 
 
 
 
 
 
 
 
 
 
 
 
 
0261f7a
40440d6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0261f7a
40440d6
 
 
 
 
 
 
 
 
 
0261f7a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# [Official] RotEx: A Rotational Expert Network for Realistic Scene Illumination Variation Synthesis

<p align="center">
    πŸ“„ <b>Paper: coming soon</b>&nbsp;&nbsp; | &nbsp;&nbsp;
    πŸ’» <a href="https://github.com/val-utehy/RotEx"><b>Source Code</b></a>&nbsp;&nbsp; | &nbsp;&nbsp;
    πŸ€— <a href="https://huggingface.co/val-utehy/RotEx"><b>Hugging Face</b></a>
</p>

<p align="center">
    <img src="demo/RotEx_Proposed.png" alt="RotEx architecture" width="1100"/>
</p>

## Introduction

**RotEx** is a transformer-based image-to-image framework for synthesizing scenes under varying illumination conditions. It combines content latents, depth information, style features, and a continuous time-of-day variable to generate illumination variations of an input scene.

### Key Features

- **Continuous illumination control:** Generate scene variations across different times of the day.
- **Illumination-dependent experts:** Adapt style and image features through time-conditioned expert routing.
- **Content and depth conditioning:** Guide synthesis using content latents and estimated scene depth.
- **Multi-step and one-step inference:** Support full-model sampling and distilled single-step generation.
- **Image and video processing:** Process image directories and videos with reference-style conditioning.

## Demo

### Illumination Variation

<p align="center">
    <img src="demo/illumination_variation_space.jpg" alt="Illumination variation" width="1000"/>
</p>

### Downstream Trajectory Comparison

<p align="center">
    <img src="demo/VLA_Traj.png" alt="AutoVLA trajectory comparison" width="1100"/>
</p>

## Installation

**Requirements:** Python β‰₯ 3.10, PyTorch β‰₯ 2.6, a compatible torchvision version, FlashAttention β‰₯ 2.0, and an NVIDIA GPU with CUDA.

```shell
git clone https://github.com/val-utehy/RotEx.git
cd RotEx

conda create -n rotex python=3.10 -y
conda activate rotex
python -m pip install --upgrade pip
```

Install CUDA-enabled PyTorch and torchvision following the [PyTorch instructions](https://pytorch.org/get-started/locally/), then run:

```shell
python -m pip install packaging ninja psutil
python -m pip install "flash-attn>=2.0.0" --no-build-isolation
python -m pip install -r requirements.txt
```

## Pretrained Models

Download weights from [Hugging Face](https://huggingface.co/val-utehy/RotEx/tree/main):

| Model | Download | Save to |
| --- | --- | --- |
| Multi-step | [best.pt](https://huggingface.co/val-utehy/RotEx/resolve/main/checkpoints/full/best.pt?download=true) | `checkpoints/full/best.pt` |
| Multi-step, final | [last.pt](https://huggingface.co/val-utehy/RotEx/resolve/main/checkpoints/full/last.pt?download=true) | `checkpoints/full/last.pt` |
| One-step | [best_distill.pt](https://huggingface.co/val-utehy/RotEx/resolve/main/distill_checkpoints/best_distill.pt?download=true) | `distill_checkpoints/best_distill.pt` |

```shell
hf download val-utehy/RotEx checkpoints/full/best.pt --local-dir .
hf download val-utehy/RotEx distill_checkpoints/best_distill.pt --local-dir .
hf download black-forest-labs/FLUX.1-dev ae.safetensors --local-dir checkpoints
```

DINOv2 and Depth Anything V2 are downloaded automatically. Run `hf auth login` if model access requires authentication.

### Expected Weights Structure

```text
RotEx/
β”œβ”€β”€ checkpoints/
β”‚   β”œβ”€β”€ ae.safetensors
β”‚   └── full/
β”‚       β”œβ”€β”€ best.pt
β”‚       └── last.pt                 # Optional
└── distill_checkpoints/
    └── best_distill.pt
```

## Usage

### 1. Training

Will be updated soon.

### 2. Multi-Step Image Inference

```shell
python infer_multistep.py \
    --checkpoint checkpoints/full/best.pt \
    --content_dir inputs/content \
    --style_dir inputs/style \
    --output_dir outputs/multistep \
    --denoise_steps 50
```

### 3. One-Step Image Inference

Set these paths in `infer_1step_style.py`:

```python
CHECKPOINT = "distill_checkpoints/best_distill.pt"
FIXED_STYLE_PATH = "inputs/style/night_reference.jpg"
DATA_PAIRS = [
    {"input": "inputs/content", "output": "outputs/1step_style"},
]
HOURS = list(range(1, 25))
```

```shell
python infer_1step_style.py
```

### 4. Multi-Step Video Inference

```shell
python infer_multistep_vid.py \
    --checkpoint checkpoints/full/best.pt \
    --video_path inputs/video.mp4 \
    --style_path inputs/style/night_reference.jpg \
    --output_dir outputs/video \
    --denoise_steps 50 \
    --hour_start 0 \
    --hour_end 24 \
    --no_audio
```

## Citation

<!-- Paper details and the official BibTeX entry will be added when available. -->

## References

- [FLUX](https://github.com/black-forest-labs/flux)
- [DINOv2](https://github.com/facebookresearch/dinov2)
- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
- [FlashAttention](https://github.com/Dao-AILab/flash-attention)

## License

This project is licensed under the [Apache License 2.0](LICENSE). Third-party components retain their respective licenses.