Spaces:
Runtime error
Runtime error
v1
Browse files- .gitattributes +1 -0
- .gitignore +3 -0
- README.md +1 -1
- app.py +218 -217
- files/videos/K_0005_IN.mp4 +3 -0
- files/videos/obama.mp4 +0 -0
- infer.py +134 -28
- pipeline.py +0 -1
- requirements.txt +3 -2
.gitattributes
CHANGED
|
@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
files/images/01.jpg filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
files/images/01.jpg filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
files/videos/K_0005_IN.mp4 filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__pycache__/
|
| 2 |
+
output/
|
| 3 |
+
gradio_cached_examples/
|
README.md
CHANGED
|
@@ -4,7 +4,7 @@ emoji: 🚀
|
|
| 4 |
colorFrom: blue
|
| 5 |
colorTo: indigo
|
| 6 |
sdk: gradio
|
| 7 |
-
sdk_version: 4.
|
| 8 |
app_file: app.py
|
| 9 |
pinned: false
|
| 10 |
license: mit
|
|
|
|
| 4 |
colorFrom: blue
|
| 5 |
colorTo: indigo
|
| 6 |
sdk: gradio
|
| 7 |
+
sdk_version: 4.44.0
|
| 8 |
app_file: app.py
|
| 9 |
pinned: false
|
| 10 |
license: mit
|
app.py
CHANGED
|
@@ -1,4 +1,3 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
from gradio_imageslider import ImageSlider
|
| 3 |
import functools
|
| 4 |
import os
|
|
@@ -14,224 +13,226 @@ from tqdm import tqdm
|
|
| 14 |
from pathlib import Path
|
| 15 |
import gradio
|
| 16 |
from gradio.utils import get_cache_folder
|
| 17 |
-
from infer import lotus
|
| 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 |
-
|
| 154 |
-
|
| 155 |
-
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
|
| 159 |
-
|
| 160 |
-
|
| 161 |
-
|
| 162 |
-
|
| 163 |
-
|
| 164 |
-
|
| 165 |
-
|
| 166 |
-
|
| 167 |
-
|
| 168 |
-
|
| 169 |
-
|
| 170 |
-
|
| 171 |
-
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
|
| 188 |
-
|
| 189 |
-
|
| 190 |
-
|
| 191 |
-
|
| 192 |
-
|
| 193 |
-
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
|
| 197 |
-
|
| 198 |
-
|
| 199 |
-
|
| 200 |
-
|
| 201 |
-
|
| 202 |
-
|
| 203 |
-
|
| 204 |
-
|
| 205 |
-
|
| 206 |
-
|
| 207 |
-
|
| 208 |
-
|
| 209 |
-
|
| 210 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 211 |
|
| 212 |
-
|
| 213 |
-
|
| 214 |
-
|
| 215 |
-
|
| 216 |
-
|
| 217 |
-
|
| 218 |
-
|
| 219 |
-
with gr.Row():
|
| 220 |
-
image_input = gr.Image()
|
| 221 |
-
image_output = gr.Image()
|
| 222 |
-
image_button = gr.Button("Flip")
|
| 223 |
|
| 224 |
-
|
| 225 |
-
|
| 226 |
-
|
| 227 |
-
|
| 228 |
-
|
| 229 |
-
|
| 230 |
-
interactive=True,
|
| 231 |
-
label="Slide me",
|
| 232 |
)
|
| 233 |
|
| 234 |
-
|
| 235 |
-
|
|
|
|
| 236 |
|
| 237 |
-
|
|
|
|
|
|
|
|
|
| 1 |
from gradio_imageslider import ImageSlider
|
| 2 |
import functools
|
| 3 |
import os
|
|
|
|
| 13 |
from pathlib import Path
|
| 14 |
import gradio
|
| 15 |
from gradio.utils import get_cache_folder
|
| 16 |
+
from infer import lotus, lotus_video
|
| 17 |
+
|
| 18 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 19 |
+
|
| 20 |
+
def infer(path_input, seed=0):
|
| 21 |
+
name_base, name_ext = os.path.splitext(os.path.basename(path_input))
|
| 22 |
+
output_g, output_d = lotus(path_input, 'depth', seed, device)
|
| 23 |
+
if not os.path.exists("files/output"):
|
| 24 |
+
os.makedirs("files/output")
|
| 25 |
+
g_save_path = os.path.join("files/output", f"{name_base}_g{name_ext}")
|
| 26 |
+
d_save_path = os.path.join("files/output", f"{name_base}_d{name_ext}")
|
| 27 |
+
output_g.save(g_save_path)
|
| 28 |
+
output_d.save(d_save_path)
|
| 29 |
+
return [path_input, g_save_path], [path_input, d_save_path]
|
| 30 |
+
|
| 31 |
+
def infer_video(path_input, seed=0):
|
| 32 |
+
frames_g, frames_d = lotus_video(path_input, 'depth', seed, device)
|
| 33 |
+
if not os.path.exists("files/output"):
|
| 34 |
+
os.makedirs("files/output")
|
| 35 |
+
name_base, _ = os.path.splitext(os.path.basename(path_input))
|
| 36 |
+
g_save_path = os.path.join("files/output", f"{name_base}_g.mp4")
|
| 37 |
+
d_save_path = os.path.join("files/output", f"{name_base}_d.mp4")
|
| 38 |
+
imageio.mimsave(g_save_path, frames_g)
|
| 39 |
+
imageio.mimsave(d_save_path, frames_d)
|
| 40 |
+
return [g_save_path, d_save_path]
|
| 41 |
+
|
| 42 |
+
def run_demo_server():
|
| 43 |
+
gradio_theme = gr.themes.Default()
|
| 44 |
+
|
| 45 |
+
with gr.Blocks(
|
| 46 |
+
theme=gradio_theme,
|
| 47 |
+
title="LOTUS (Depth)",
|
| 48 |
+
css="""
|
| 49 |
+
#download {
|
| 50 |
+
height: 118px;
|
| 51 |
+
}
|
| 52 |
+
.slider .inner {
|
| 53 |
+
width: 5px;
|
| 54 |
+
background: #FFF;
|
| 55 |
+
}
|
| 56 |
+
.viewport {
|
| 57 |
+
aspect-ratio: 4/3;
|
| 58 |
+
}
|
| 59 |
+
.tabs button.selected {
|
| 60 |
+
font-size: 20px !important;
|
| 61 |
+
color: crimson !important;
|
| 62 |
+
}
|
| 63 |
+
h1 {
|
| 64 |
+
text-align: center;
|
| 65 |
+
display: block;
|
| 66 |
+
}
|
| 67 |
+
h2 {
|
| 68 |
+
text-align: center;
|
| 69 |
+
display: block;
|
| 70 |
+
}
|
| 71 |
+
h3 {
|
| 72 |
+
text-align: center;
|
| 73 |
+
display: block;
|
| 74 |
+
}
|
| 75 |
+
.md_feedback li {
|
| 76 |
+
margin-bottom: 0px !important;
|
| 77 |
+
}
|
| 78 |
+
""",
|
| 79 |
+
head="""
|
| 80 |
+
<script async src="https://www.googletagmanager.com/gtag/js?id=G-1FWSVCGZTG"></script>
|
| 81 |
+
<script>
|
| 82 |
+
window.dataLayer = window.dataLayer || [];
|
| 83 |
+
function gtag() {dataLayer.push(arguments);}
|
| 84 |
+
gtag('js', new Date());
|
| 85 |
+
gtag('config', 'G-1FWSVCGZTG');
|
| 86 |
+
</script>
|
| 87 |
+
""",
|
| 88 |
+
) as demo:
|
| 89 |
+
gr.Markdown(
|
| 90 |
+
"""
|
| 91 |
+
# LOTUS: Diffusion-based Visual Foundation Model for High-quality Dense Prediction
|
| 92 |
+
<p align="center">
|
| 93 |
+
<a title="Page" href="https://lotus3d.github.io/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
|
| 94 |
+
<img src="https://img.shields.io/badge/Project-Website-pink?logo=googlechrome&logoColor=white">
|
| 95 |
+
</a>
|
| 96 |
+
<a title="arXiv" href="https://arxiv.org/abs/2409.18124" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
|
| 97 |
+
<img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white">
|
| 98 |
+
</a>
|
| 99 |
+
<a title="Github" href="https://github.com/EnVision-Research/Lotus" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
|
| 100 |
+
<img src="https://img.shields.io/github/stars/EnVision-Research/Lotus?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
|
| 101 |
+
</a>
|
| 102 |
+
<a title="Social" href="https://x.com/haodongli00/status/1839524569058582884" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
|
| 103 |
+
<img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
|
| 104 |
+
</a>
|
| 105 |
+
"""
|
| 106 |
+
)
|
| 107 |
+
with gr.Tabs(elem_classes=["tabs"]):
|
| 108 |
+
with gr.Tab("IMAGE"):
|
| 109 |
+
with gr.Row():
|
| 110 |
+
with gr.Column():
|
| 111 |
+
image_input = gr.Image(
|
| 112 |
+
label="Input Image",
|
| 113 |
+
type="filepath",
|
| 114 |
+
)
|
| 115 |
+
seed = gr.Number(
|
| 116 |
+
label="Seed (only for Generative mode)",
|
| 117 |
+
minimum=0,
|
| 118 |
+
maximum=999999999,
|
| 119 |
+
)
|
| 120 |
+
with gr.Row():
|
| 121 |
+
image_submit_btn = gr.Button(
|
| 122 |
+
value="Predict Depth!", variant="primary"
|
| 123 |
+
)
|
| 124 |
+
image_reset_btn = gr.Button(value="Reset")
|
| 125 |
+
with gr.Column():
|
| 126 |
+
image_output_g = ImageSlider(
|
| 127 |
+
label="Output (Generative)",
|
| 128 |
+
type="filepath",
|
| 129 |
+
interactive=False,
|
| 130 |
+
elem_classes="slider",
|
| 131 |
+
position=0.25,
|
| 132 |
+
)
|
| 133 |
+
with gr.Row():
|
| 134 |
+
image_output_d = ImageSlider(
|
| 135 |
+
label="Output (Discriminative)",
|
| 136 |
+
type="filepath",
|
| 137 |
+
interactive=False,
|
| 138 |
+
elem_classes="slider",
|
| 139 |
+
position=0.25,
|
| 140 |
+
)
|
| 141 |
+
|
| 142 |
+
gr.Examples(
|
| 143 |
+
fn=infer,
|
| 144 |
+
examples=sorted([
|
| 145 |
+
os.path.join("files", "images", name)
|
| 146 |
+
for name in os.listdir(os.path.join("files", "images"))
|
| 147 |
+
]),
|
| 148 |
+
inputs=[image_input],
|
| 149 |
+
outputs=[image_output_g, image_output_d],
|
| 150 |
+
cache_examples=True,
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
with gr.Tab("VIDEO"):
|
| 154 |
+
with gr.Row():
|
| 155 |
+
with gr.Column():
|
| 156 |
+
input_video = gr.Video(
|
| 157 |
+
label="Input Video",
|
| 158 |
+
autoplay=True,
|
| 159 |
+
loop=True,
|
| 160 |
+
)
|
| 161 |
+
seed = gr.Number(
|
| 162 |
+
label="Seed (only for Generative mode)",
|
| 163 |
+
minimum=0,
|
| 164 |
+
maximum=999999999,
|
| 165 |
+
)
|
| 166 |
+
with gr.Row():
|
| 167 |
+
video_submit_btn = gr.Button(
|
| 168 |
+
value="Compute Depth!", variant="primary"
|
| 169 |
+
)
|
| 170 |
+
video_reset_btn = gr.Button(value="Reset")
|
| 171 |
+
with gr.Column():
|
| 172 |
+
video_output_g = gr.Video(
|
| 173 |
+
label="Output (Generative)",
|
| 174 |
+
interactive=False,
|
| 175 |
+
autoplay=True,
|
| 176 |
+
loop=True,
|
| 177 |
+
show_share_button=True,
|
| 178 |
+
)
|
| 179 |
+
with gr.Row():
|
| 180 |
+
video_output_d = gr.Video(
|
| 181 |
+
label="Output (Discriminative)",
|
| 182 |
+
interactive=False,
|
| 183 |
+
autoplay=True,
|
| 184 |
+
loop=True,
|
| 185 |
+
show_share_button=True,
|
| 186 |
+
)
|
| 187 |
+
|
| 188 |
+
gr.Examples(
|
| 189 |
+
fn=infer_video,
|
| 190 |
+
examples=sorted([
|
| 191 |
+
os.path.join("files", "videos", name)
|
| 192 |
+
for name in os.listdir(os.path.join("files", "videos"))
|
| 193 |
+
]),
|
| 194 |
+
inputs=[input_video],
|
| 195 |
+
outputs=[video_output_g, video_output_d],
|
| 196 |
+
cache_examples=True,
|
| 197 |
+
)
|
| 198 |
+
|
| 199 |
+
### Image
|
| 200 |
+
image_submit_btn.click(
|
| 201 |
+
fn=infer,
|
| 202 |
+
inputs=[image_input, seed],
|
| 203 |
+
outputs=[image_output_g, image_output_d],
|
| 204 |
+
concurrency_limit=1,
|
| 205 |
+
)
|
| 206 |
+
image_reset_btn.click(
|
| 207 |
+
fn=lambda: (
|
| 208 |
+
None,
|
| 209 |
+
None,
|
| 210 |
+
None,
|
| 211 |
+
),
|
| 212 |
+
inputs=[],
|
| 213 |
+
outputs=[image_output_g, image_output_d],
|
| 214 |
+
queue=False,
|
| 215 |
+
)
|
| 216 |
|
| 217 |
+
### Video
|
| 218 |
+
video_submit_btn.click(
|
| 219 |
+
fn=infer_video,
|
| 220 |
+
inputs=[input_video, seed],
|
| 221 |
+
outputs=[video_output_g, video_output_d],
|
| 222 |
+
queue=True,
|
| 223 |
+
)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 224 |
|
| 225 |
+
### Server launch
|
| 226 |
+
demo.queue(
|
| 227 |
+
api_open=False,
|
| 228 |
+
).launch(
|
| 229 |
+
server_name="0.0.0.0",
|
| 230 |
+
server_port=7860,
|
|
|
|
|
|
|
| 231 |
)
|
| 232 |
|
| 233 |
+
def main():
|
| 234 |
+
os.system("pip freeze")
|
| 235 |
+
run_demo_server()
|
| 236 |
|
| 237 |
+
if __name__ == "__main__":
|
| 238 |
+
main()
|
files/videos/K_0005_IN.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9a532ba2738716dbb244e0d7172cf681879218cbbdad09980404fa08ef6b9ecc
|
| 3 |
+
size 3095352
|
files/videos/obama.mp4
DELETED
|
Binary file (320 kB)
|
|
|
infer.py
CHANGED
|
@@ -14,6 +14,9 @@ from pipeline import LotusGPipeline, LotusDPipeline
|
|
| 14 |
from utils.image_utils import colorize_depth_map
|
| 15 |
from utils.seed_all import seed_all
|
| 16 |
|
|
|
|
|
|
|
|
|
|
| 17 |
check_min_version('0.28.0.dev0')
|
| 18 |
|
| 19 |
def infer_pipe(pipe, image_input, task_name, seed, device):
|
|
@@ -22,36 +25,137 @@ def infer_pipe(pipe, image_input, task_name, seed, device):
|
|
| 22 |
else:
|
| 23 |
generator = torch.Generator(device=device).manual_seed(seed)
|
| 24 |
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
|
| 31 |
-
|
| 32 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
|
| 34 |
-
|
| 35 |
-
pred = pipe(
|
| 36 |
-
rgb_in=test_image,
|
| 37 |
-
prompt='',
|
| 38 |
-
num_inference_steps=1,
|
| 39 |
-
generator=generator,
|
| 40 |
-
# guidance_scale=0,
|
| 41 |
-
output_type='np',
|
| 42 |
-
timesteps=[999],
|
| 43 |
-
task_emb=task_emb,
|
| 44 |
-
).images[0]
|
| 45 |
-
|
| 46 |
-
# Post-process the prediction
|
| 47 |
if task_name == 'depth':
|
| 48 |
-
|
| 49 |
-
|
| 50 |
else:
|
| 51 |
-
|
| 52 |
-
|
| 53 |
|
| 54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
|
| 56 |
def lotus(image_input, task_name, seed, device):
|
| 57 |
if task_name == 'depth':
|
|
@@ -61,7 +165,7 @@ def lotus(image_input, task_name, seed, device):
|
|
| 61 |
model_g = 'jingheya/lotus-normal-g-v1-0'
|
| 62 |
model_d = 'jingheya/lotus-normal-d-v1-0'
|
| 63 |
|
| 64 |
-
dtype = torch.
|
| 65 |
pipe_g = LotusGPipeline.from_pretrained(
|
| 66 |
model_g,
|
| 67 |
torch_dtype=dtype,
|
|
@@ -72,6 +176,8 @@ def lotus(image_input, task_name, seed, device):
|
|
| 72 |
)
|
| 73 |
pipe_g.to(device)
|
| 74 |
pipe_d.to(device)
|
|
|
|
|
|
|
| 75 |
logging.info(f"Successfully loading pipeline from {model_g} and {model_d}.")
|
| 76 |
output_g = infer_pipe(pipe_g, image_input, task_name, seed, device)
|
| 77 |
output_d = infer_pipe(pipe_d, image_input, task_name, seed, device)
|
|
@@ -158,7 +264,7 @@ def main():
|
|
| 158 |
dtype = torch.float16
|
| 159 |
logging.info(f"Running with half precision ({dtype}).")
|
| 160 |
else:
|
| 161 |
-
dtype = torch.
|
| 162 |
|
| 163 |
# -------------------- Device --------------------
|
| 164 |
if torch.cuda.is_available():
|
|
@@ -206,7 +312,7 @@ def main():
|
|
| 206 |
for i in tqdm(range(len(test_images))):
|
| 207 |
# Preprocess validation image
|
| 208 |
test_image = Image.open(test_images[i]).convert('RGB')
|
| 209 |
-
test_image = np.array(test_image).astype(np.
|
| 210 |
test_image = torch.tensor(test_image).permute(2,0,1).unsqueeze(0)
|
| 211 |
test_image = test_image / 127.5 - 1.0
|
| 212 |
test_image = test_image.to(device)
|
|
|
|
| 14 |
from utils.image_utils import colorize_depth_map
|
| 15 |
from utils.seed_all import seed_all
|
| 16 |
|
| 17 |
+
from contextlib import nullcontext
|
| 18 |
+
import cv2
|
| 19 |
+
|
| 20 |
check_min_version('0.28.0.dev0')
|
| 21 |
|
| 22 |
def infer_pipe(pipe, image_input, task_name, seed, device):
|
|
|
|
| 25 |
else:
|
| 26 |
generator = torch.Generator(device=device).manual_seed(seed)
|
| 27 |
|
| 28 |
+
if torch.backends.mps.is_available():
|
| 29 |
+
autocast_ctx = nullcontext()
|
| 30 |
+
else:
|
| 31 |
+
autocast_ctx = torch.autocast(pipe.device.type)
|
| 32 |
+
with autocast_ctx:
|
| 33 |
|
| 34 |
+
test_image = Image.open(image_input).convert('RGB')
|
| 35 |
+
test_image = np.array(test_image).astype(np.float16)
|
| 36 |
+
test_image = torch.tensor(test_image).permute(2,0,1).unsqueeze(0)
|
| 37 |
+
test_image = test_image / 127.5 - 1.0
|
| 38 |
+
test_image = test_image.to(device)
|
| 39 |
+
|
| 40 |
+
task_emb = torch.tensor([1, 0]).float().unsqueeze(0).repeat(1, 1).to(device)
|
| 41 |
+
task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1).repeat(1, 1)
|
| 42 |
+
|
| 43 |
+
# Run
|
| 44 |
+
pred = pipe(
|
| 45 |
+
rgb_in=test_image,
|
| 46 |
+
prompt='',
|
| 47 |
+
num_inference_steps=1,
|
| 48 |
+
generator=generator,
|
| 49 |
+
# guidance_scale=0,
|
| 50 |
+
output_type='np',
|
| 51 |
+
timesteps=[999],
|
| 52 |
+
task_emb=task_emb,
|
| 53 |
+
).images[0]
|
| 54 |
+
|
| 55 |
+
# Post-process the prediction
|
| 56 |
+
if task_name == 'depth':
|
| 57 |
+
output_npy = pred.mean(axis=-1)
|
| 58 |
+
output_color = colorize_depth_map(output_npy)
|
| 59 |
+
else:
|
| 60 |
+
output_npy = pred
|
| 61 |
+
output_color = Image.fromarray((output_npy * 255).astype(np.uint8))
|
| 62 |
+
|
| 63 |
+
return output_color
|
| 64 |
|
| 65 |
+
def lotus_video(input_video, task_name, seed, device):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 66 |
if task_name == 'depth':
|
| 67 |
+
model_g = 'jingheya/lotus-depth-g-v1-0'
|
| 68 |
+
model_d = 'jingheya/lotus-depth-d-v1-0'
|
| 69 |
else:
|
| 70 |
+
model_g = 'jingheya/lotus-normal-g-v1-0'
|
| 71 |
+
model_d = 'jingheya/lotus-normal-d-v1-0'
|
| 72 |
|
| 73 |
+
dtype = torch.float16
|
| 74 |
+
pipe_g = LotusGPipeline.from_pretrained(
|
| 75 |
+
model_g,
|
| 76 |
+
torch_dtype=dtype,
|
| 77 |
+
)
|
| 78 |
+
pipe_d = LotusDPipeline.from_pretrained(
|
| 79 |
+
model_d,
|
| 80 |
+
torch_dtype=dtype,
|
| 81 |
+
)
|
| 82 |
+
pipe_g.to(device)
|
| 83 |
+
pipe_d.to(device)
|
| 84 |
+
pipe_g.set_progress_bar_config(disable=True)
|
| 85 |
+
pipe_d.set_progress_bar_config(disable=True)
|
| 86 |
+
logging.info(f"Successfully loading pipeline from {model_g} and {model_d}.")
|
| 87 |
+
|
| 88 |
+
# load the video and split it into frames
|
| 89 |
+
cap = cv2.VideoCapture(input_video)
|
| 90 |
+
frames = []
|
| 91 |
+
while True:
|
| 92 |
+
ret, frame = cap.read()
|
| 93 |
+
if not ret:
|
| 94 |
+
break
|
| 95 |
+
frames.append(frame)
|
| 96 |
+
cap.release()
|
| 97 |
+
logging.info(f"There are {len(frames)} frames in the video.")
|
| 98 |
+
|
| 99 |
+
if seed is None:
|
| 100 |
+
generator = None
|
| 101 |
+
else:
|
| 102 |
+
generator = torch.Generator(device=device).manual_seed(seed)
|
| 103 |
+
|
| 104 |
+
task_emb = torch.tensor([1, 0]).float().unsqueeze(0).repeat(1, 1).to(device)
|
| 105 |
+
task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1).repeat(1, 1)
|
| 106 |
+
|
| 107 |
+
output_g = []
|
| 108 |
+
output_d = []
|
| 109 |
+
for frame in frames:
|
| 110 |
+
if torch.backends.mps.is_available():
|
| 111 |
+
autocast_ctx = nullcontext()
|
| 112 |
+
else:
|
| 113 |
+
autocast_ctx = torch.autocast(pipe_g.device.type)
|
| 114 |
+
with autocast_ctx:
|
| 115 |
+
test_image = frame
|
| 116 |
+
test_image = np.array(test_image).astype(np.float16)
|
| 117 |
+
test_image = torch.tensor(test_image).permute(2,0,1).unsqueeze(0)
|
| 118 |
+
test_image = test_image / 127.5 - 1.0
|
| 119 |
+
test_image = test_image.to(device)
|
| 120 |
+
|
| 121 |
+
# Run
|
| 122 |
+
pred_g = pipe_g(
|
| 123 |
+
rgb_in=test_image,
|
| 124 |
+
prompt='',
|
| 125 |
+
num_inference_steps=1,
|
| 126 |
+
generator=generator,
|
| 127 |
+
# guidance_scale=0,
|
| 128 |
+
output_type='np',
|
| 129 |
+
timesteps=[999],
|
| 130 |
+
task_emb=task_emb,
|
| 131 |
+
).images[0]
|
| 132 |
+
pred_d = pipe_d(
|
| 133 |
+
rgb_in=test_image,
|
| 134 |
+
prompt='',
|
| 135 |
+
num_inference_steps=1,
|
| 136 |
+
generator=generator,
|
| 137 |
+
# guidance_scale=0,
|
| 138 |
+
output_type='np',
|
| 139 |
+
timesteps=[999],
|
| 140 |
+
task_emb=task_emb,
|
| 141 |
+
).images[0]
|
| 142 |
+
|
| 143 |
+
# Post-process the prediction
|
| 144 |
+
if task_name == 'depth':
|
| 145 |
+
output_npy_g = pred_g.mean(axis=-1)
|
| 146 |
+
output_color_g = colorize_depth_map(output_npy_g)
|
| 147 |
+
output_npy_d = pred_d.mean(axis=-1)
|
| 148 |
+
output_color_d = colorize_depth_map(output_npy_d)
|
| 149 |
+
else:
|
| 150 |
+
output_npy_g = pred_g
|
| 151 |
+
output_color_g = Image.fromarray((output_npy_g * 255).astype(np.uint8))
|
| 152 |
+
output_npy_d = pred_d
|
| 153 |
+
output_color_d = Image.fromarray((output_npy_d * 255).astype(np.uint8))
|
| 154 |
+
|
| 155 |
+
output_g.append(output_color_g)
|
| 156 |
+
output_d.append(output_color_d)
|
| 157 |
+
|
| 158 |
+
return output_g, output_d
|
| 159 |
|
| 160 |
def lotus(image_input, task_name, seed, device):
|
| 161 |
if task_name == 'depth':
|
|
|
|
| 165 |
model_g = 'jingheya/lotus-normal-g-v1-0'
|
| 166 |
model_d = 'jingheya/lotus-normal-d-v1-0'
|
| 167 |
|
| 168 |
+
dtype = torch.float16
|
| 169 |
pipe_g = LotusGPipeline.from_pretrained(
|
| 170 |
model_g,
|
| 171 |
torch_dtype=dtype,
|
|
|
|
| 176 |
)
|
| 177 |
pipe_g.to(device)
|
| 178 |
pipe_d.to(device)
|
| 179 |
+
pipe_g.set_progress_bar_config(disable=True)
|
| 180 |
+
pipe_d.set_progress_bar_config(disable=True)
|
| 181 |
logging.info(f"Successfully loading pipeline from {model_g} and {model_d}.")
|
| 182 |
output_g = infer_pipe(pipe_g, image_input, task_name, seed, device)
|
| 183 |
output_d = infer_pipe(pipe_d, image_input, task_name, seed, device)
|
|
|
|
| 264 |
dtype = torch.float16
|
| 265 |
logging.info(f"Running with half precision ({dtype}).")
|
| 266 |
else:
|
| 267 |
+
dtype = torch.float16
|
| 268 |
|
| 269 |
# -------------------- Device --------------------
|
| 270 |
if torch.cuda.is_available():
|
|
|
|
| 312 |
for i in tqdm(range(len(test_images))):
|
| 313 |
# Preprocess validation image
|
| 314 |
test_image = Image.open(test_images[i]).convert('RGB')
|
| 315 |
+
test_image = np.array(test_image).astype(np.float16)
|
| 316 |
test_image = torch.tensor(test_image).permute(2,0,1).unsqueeze(0)
|
| 317 |
test_image = test_image / 127.5 - 1.0
|
| 318 |
test_image = test_image.to(device)
|
pipeline.py
CHANGED
|
@@ -1197,7 +1197,6 @@ class LotusGPipeline(DirectDiffusionPipeline):
|
|
| 1197 |
# 2. Define call parameters
|
| 1198 |
batch_size = rgb_in.shape[0]
|
| 1199 |
device = self._execution_device
|
| 1200 |
-
print("Device: ", device)
|
| 1201 |
|
| 1202 |
# 3. Encode input prompt
|
| 1203 |
prompt_embeds, _ = self.encode_prompt(
|
|
|
|
| 1197 |
# 2. Define call parameters
|
| 1198 |
batch_size = rgb_in.shape[0]
|
| 1199 |
device = self._execution_device
|
|
|
|
| 1200 |
|
| 1201 |
# 3. Encode input prompt
|
| 1202 |
prompt_embeds, _ = self.encode_prompt(
|
requirements.txt
CHANGED
|
@@ -17,7 +17,8 @@ h5py==3.11.0
|
|
| 17 |
omegaconf==2.3.0
|
| 18 |
tabulate==0.9.0
|
| 19 |
imageio==2.35.1
|
|
|
|
| 20 |
spaces==0.28.3
|
| 21 |
-
gradio==4.
|
| 22 |
gradio-imageslider==0.0.16
|
| 23 |
-
|
|
|
|
| 17 |
omegaconf==2.3.0
|
| 18 |
tabulate==0.9.0
|
| 19 |
imageio==2.35.1
|
| 20 |
+
imageio-ffmpeg==0.5.1
|
| 21 |
spaces==0.28.3
|
| 22 |
+
gradio==4.44.0
|
| 23 |
gradio-imageslider==0.0.16
|
| 24 |
+
gradio-client==1.3.0
|