lemonaddie
commited on
Commit
•
8e0bf0b
1
Parent(s):
671d862
Upload app.py
Browse files
app.py
ADDED
@@ -0,0 +1,475 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import functools
|
2 |
+
import os
|
3 |
+
import shutil
|
4 |
+
import sys
|
5 |
+
|
6 |
+
import git
|
7 |
+
import gradio as gr
|
8 |
+
import numpy as np
|
9 |
+
import torch as torch
|
10 |
+
from PIL import Image
|
11 |
+
|
12 |
+
from gradio_imageslider import ImageSlider
|
13 |
+
|
14 |
+
from extrude import extrude_depth_3d
|
15 |
+
|
16 |
+
|
17 |
+
def process(
|
18 |
+
pipe,
|
19 |
+
path_input,
|
20 |
+
ensemble_size,
|
21 |
+
denoise_steps,
|
22 |
+
processing_res,
|
23 |
+
path_out_16bit=None,
|
24 |
+
path_out_fp32=None,
|
25 |
+
path_out_vis=None,
|
26 |
+
_input_3d_plane_near=None,
|
27 |
+
_input_3d_plane_far=None,
|
28 |
+
_input_3d_embossing=None,
|
29 |
+
_input_3d_filter_size=None,
|
30 |
+
_input_3d_frame_near=None,
|
31 |
+
):
|
32 |
+
if path_out_vis is not None:
|
33 |
+
return (
|
34 |
+
[path_out_16bit, path_out_vis],
|
35 |
+
[path_out_16bit, path_out_fp32, path_out_vis],
|
36 |
+
)
|
37 |
+
|
38 |
+
input_image = Image.open(path_input)
|
39 |
+
|
40 |
+
pipe_out = pipe(
|
41 |
+
input_image,
|
42 |
+
ensemble_size=ensemble_size,
|
43 |
+
denoising_steps=denoise_steps,
|
44 |
+
processing_res=processing_res,
|
45 |
+
batch_size=1 if processing_res == 0 else 0,
|
46 |
+
show_progress_bar=True,
|
47 |
+
)
|
48 |
+
|
49 |
+
depth_pred = pipe_out.depth_np
|
50 |
+
depth_colored = pipe_out.depth_colored
|
51 |
+
depth_16bit = (depth_pred * 65535.0).astype(np.uint16)
|
52 |
+
|
53 |
+
path_output_dir = os.path.splitext(path_input)[0] + "_output"
|
54 |
+
os.makedirs(path_output_dir, exist_ok=True)
|
55 |
+
|
56 |
+
name_base = os.path.splitext(os.path.basename(path_input))[0]
|
57 |
+
path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy")
|
58 |
+
path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png")
|
59 |
+
path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png")
|
60 |
+
|
61 |
+
np.save(path_out_fp32, depth_pred)
|
62 |
+
Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16")
|
63 |
+
depth_colored.save(path_out_vis)
|
64 |
+
|
65 |
+
return (
|
66 |
+
[path_out_16bit, path_out_vis],
|
67 |
+
[path_out_16bit, path_out_fp32, path_out_vis],
|
68 |
+
)
|
69 |
+
|
70 |
+
|
71 |
+
def process_3d(
|
72 |
+
input_image,
|
73 |
+
files,
|
74 |
+
size_longest_px,
|
75 |
+
size_longest_cm,
|
76 |
+
filter_size,
|
77 |
+
plane_near,
|
78 |
+
plane_far,
|
79 |
+
embossing,
|
80 |
+
frame_thickness,
|
81 |
+
frame_near,
|
82 |
+
frame_far,
|
83 |
+
):
|
84 |
+
if input_image is None or len(files) < 1:
|
85 |
+
raise gr.Error("Please upload an image (or use examples) and compute depth first")
|
86 |
+
|
87 |
+
if plane_near >= plane_far:
|
88 |
+
raise gr.Error("NEAR plane must have a value smaller than the FAR plane")
|
89 |
+
|
90 |
+
def _process_3d(size_longest_px, filter_size, vertex_colors, scene_lights, output_model_scale=None):
|
91 |
+
image_rgb = input_image
|
92 |
+
image_depth = files[0]
|
93 |
+
|
94 |
+
image_rgb_basename, image_rgb_ext = os.path.splitext(image_rgb)
|
95 |
+
image_depth_basename, image_depth_ext = os.path.splitext(image_depth)
|
96 |
+
|
97 |
+
image_rgb_content = Image.open(image_rgb)
|
98 |
+
image_rgb_w, image_rgb_h = image_rgb_content.width, image_rgb_content.height
|
99 |
+
image_rgb_d = max(image_rgb_w, image_rgb_h)
|
100 |
+
image_new_w = size_longest_px * image_rgb_w // image_rgb_d
|
101 |
+
image_new_h = size_longest_px * image_rgb_h // image_rgb_d
|
102 |
+
|
103 |
+
image_rgb_new = image_rgb_basename + f"_{size_longest_px}" + image_rgb_ext
|
104 |
+
image_depth_new = image_depth_basename + f"_{size_longest_px}" + image_depth_ext
|
105 |
+
image_rgb_content.resize((image_new_w, image_new_h), Image.LANCZOS).save(
|
106 |
+
image_rgb_new
|
107 |
+
)
|
108 |
+
Image.open(image_depth).resize((image_new_w, image_new_h), Image.LANCZOS).save(
|
109 |
+
image_depth_new
|
110 |
+
)
|
111 |
+
|
112 |
+
path_glb, path_stl = extrude_depth_3d(
|
113 |
+
image_rgb_new,
|
114 |
+
image_depth_new,
|
115 |
+
output_model_scale=size_longest_cm * 10 if output_model_scale is None else output_model_scale,
|
116 |
+
filter_size=filter_size,
|
117 |
+
coef_near=plane_near,
|
118 |
+
coef_far=plane_far,
|
119 |
+
emboss=embossing / 100,
|
120 |
+
f_thic=frame_thickness / 100,
|
121 |
+
f_near=frame_near / 100,
|
122 |
+
f_back=frame_far / 100,
|
123 |
+
vertex_colors=vertex_colors,
|
124 |
+
scene_lights=scene_lights,
|
125 |
+
)
|
126 |
+
|
127 |
+
return path_glb, path_stl
|
128 |
+
|
129 |
+
path_viewer_glb, _ = _process_3d(256, filter_size, vertex_colors=False, scene_lights=True, output_model_scale=1)
|
130 |
+
path_files_glb, path_files_stl = _process_3d(size_longest_px, filter_size, vertex_colors=True, scene_lights=False)
|
131 |
+
|
132 |
+
# sanitize 3d viewer glb path to keep babylon.js happy
|
133 |
+
path_viewer_glb_sanitized = os.path.join(os.path.dirname(path_viewer_glb), "preview.glb")
|
134 |
+
if path_viewer_glb_sanitized != path_viewer_glb:
|
135 |
+
os.rename(path_viewer_glb, path_viewer_glb_sanitized)
|
136 |
+
path_viewer_glb = path_viewer_glb_sanitized
|
137 |
+
|
138 |
+
return path_viewer_glb, [path_files_glb, path_files_stl]
|
139 |
+
|
140 |
+
|
141 |
+
def run_demo_server(pipe):
|
142 |
+
process_pipe = functools.partial(process, pipe)
|
143 |
+
os.environ["GRADIO_ALLOW_FLAGGING"] = "never"
|
144 |
+
|
145 |
+
with gr.Blocks(
|
146 |
+
analytics_enabled=False,
|
147 |
+
title="Geowizard Depth and Normal Estimation",
|
148 |
+
css="""
|
149 |
+
#download {
|
150 |
+
height: 118px;
|
151 |
+
}
|
152 |
+
.slider .inner {
|
153 |
+
width: 5px;
|
154 |
+
background: #FFF;
|
155 |
+
}
|
156 |
+
.viewport {
|
157 |
+
aspect-ratio: 4/3;
|
158 |
+
}
|
159 |
+
""",
|
160 |
+
) as demo:
|
161 |
+
gr.Markdown(
|
162 |
+
"""
|
163 |
+
"""
|
164 |
+
)
|
165 |
+
|
166 |
+
with gr.Row():
|
167 |
+
with gr.Column():
|
168 |
+
input_image = gr.Image(
|
169 |
+
label="Input Image",
|
170 |
+
type="filepath",
|
171 |
+
)
|
172 |
+
with gr.Accordion("Advanced options", open=False):
|
173 |
+
ensemble_size = gr.Slider(
|
174 |
+
label="Ensemble size",
|
175 |
+
minimum=1,
|
176 |
+
maximum=20,
|
177 |
+
step=1,
|
178 |
+
value=10,
|
179 |
+
)
|
180 |
+
denoise_steps = gr.Slider(
|
181 |
+
label="Number of denoising steps",
|
182 |
+
minimum=1,
|
183 |
+
maximum=20,
|
184 |
+
step=1,
|
185 |
+
value=10,
|
186 |
+
)
|
187 |
+
processing_res = gr.Radio(
|
188 |
+
[
|
189 |
+
("Native", 0),
|
190 |
+
("Recommended", 768),
|
191 |
+
],
|
192 |
+
label="Processing resolution",
|
193 |
+
value=768,
|
194 |
+
)
|
195 |
+
input_output_16bit = gr.File(
|
196 |
+
label="Predicted depth (16-bit)",
|
197 |
+
visible=False,
|
198 |
+
)
|
199 |
+
input_output_fp32 = gr.File(
|
200 |
+
label="Predicted depth (32-bit)",
|
201 |
+
visible=False,
|
202 |
+
)
|
203 |
+
input_output_vis = gr.File(
|
204 |
+
label="Predicted depth (red-near, blue-far)",
|
205 |
+
visible=False,
|
206 |
+
)
|
207 |
+
with gr.Row():
|
208 |
+
submit_btn = gr.Button(value="Compute Depth", variant="primary")
|
209 |
+
clear_btn = gr.Button(value="Clear")
|
210 |
+
with gr.Column():
|
211 |
+
output_slider = ImageSlider(
|
212 |
+
label="Predicted depth (red-near, blue-far)",
|
213 |
+
type="filepath",
|
214 |
+
show_download_button=True,
|
215 |
+
show_share_button=True,
|
216 |
+
interactive=False,
|
217 |
+
elem_classes="slider",
|
218 |
+
position=0.25,
|
219 |
+
)
|
220 |
+
files = gr.Files(
|
221 |
+
label="Depth outputs",
|
222 |
+
elem_id="download",
|
223 |
+
interactive=False,
|
224 |
+
)
|
225 |
+
|
226 |
+
demo_3d_header = gr.Markdown(
|
227 |
+
"""
|
228 |
+
<h3 align="center">Depth Maps</h3>
|
229 |
+
<p align="justify">
|
230 |
+
TBD
|
231 |
+
result (see Pro Tips below).
|
232 |
+
</p>
|
233 |
+
""",
|
234 |
+
render=False,
|
235 |
+
)
|
236 |
+
|
237 |
+
demo_3d = gr.Row(render=False)
|
238 |
+
with demo_3d:
|
239 |
+
with gr.Column():
|
240 |
+
with gr.Accordion("3D printing demo: Main options", open=True):
|
241 |
+
plane_near = gr.Slider(
|
242 |
+
label="Relative position of the near plane (between 0 and 1)",
|
243 |
+
minimum=0.0,
|
244 |
+
maximum=1.0,
|
245 |
+
step=0.001,
|
246 |
+
value=0.0,
|
247 |
+
)
|
248 |
+
plane_far = gr.Slider(
|
249 |
+
label="Relative position of the far plane (between near and 1)",
|
250 |
+
minimum=0.0,
|
251 |
+
maximum=1.0,
|
252 |
+
step=0.001,
|
253 |
+
value=1.0,
|
254 |
+
)
|
255 |
+
embossing = gr.Slider(
|
256 |
+
label="Embossing level",
|
257 |
+
minimum=0,
|
258 |
+
maximum=100,
|
259 |
+
step=1,
|
260 |
+
value=20,
|
261 |
+
)
|
262 |
+
with gr.Accordion("3D printing demo: Advanced options", open=False):
|
263 |
+
size_longest_px = gr.Slider(
|
264 |
+
label="Size (px) of the longest side",
|
265 |
+
minimum=256,
|
266 |
+
maximum=1024,
|
267 |
+
step=256,
|
268 |
+
value=512,
|
269 |
+
)
|
270 |
+
size_longest_cm = gr.Slider(
|
271 |
+
label="Size (cm) of the longest side",
|
272 |
+
minimum=1,
|
273 |
+
maximum=100,
|
274 |
+
step=1,
|
275 |
+
value=10,
|
276 |
+
)
|
277 |
+
filter_size = gr.Slider(
|
278 |
+
label="Size (px) of the smoothing filter",
|
279 |
+
minimum=1,
|
280 |
+
maximum=5,
|
281 |
+
step=2,
|
282 |
+
value=3,
|
283 |
+
)
|
284 |
+
frame_thickness = gr.Slider(
|
285 |
+
label="Frame thickness",
|
286 |
+
minimum=0,
|
287 |
+
maximum=100,
|
288 |
+
step=1,
|
289 |
+
value=5,
|
290 |
+
)
|
291 |
+
frame_near = gr.Slider(
|
292 |
+
label="Frame's near plane offset",
|
293 |
+
minimum=-100,
|
294 |
+
maximum=100,
|
295 |
+
step=1,
|
296 |
+
value=1,
|
297 |
+
)
|
298 |
+
frame_far = gr.Slider(
|
299 |
+
label="Frame's far plane offset",
|
300 |
+
minimum=1,
|
301 |
+
maximum=10,
|
302 |
+
step=1,
|
303 |
+
value=1,
|
304 |
+
)
|
305 |
+
with gr.Row():
|
306 |
+
submit_3d = gr.Button(value="Create 3D", variant="primary")
|
307 |
+
clear_3d = gr.Button(value="Clear 3D")
|
308 |
+
gr.Markdown(
|
309 |
+
"""
|
310 |
+
<h5 align="center">Pro Tips</h5>
|
311 |
+
<ol>
|
312 |
+
<li><b>Re-render with new parameters</b>: Click "Clear 3D" and then "Create 3D".</li>
|
313 |
+
<li><b>Adjust 3D scale and cut-off focus</b>: Set the frame's near plane offset to the
|
314 |
+
minimum and use 3D preview to evaluate depth scaling. Repeat until the scale is correct and
|
315 |
+
everything important is in the focus. Set the optimal value for frame's near
|
316 |
+
plane offset as a last step.</li>
|
317 |
+
<li><b>Increase details</b>: Decrease size of the smoothing filter (also increases noise).</li>
|
318 |
+
</ol>
|
319 |
+
"""
|
320 |
+
)
|
321 |
+
|
322 |
+
with gr.Column():
|
323 |
+
viewer_3d = gr.Model3D(
|
324 |
+
camera_position=(75.0, 90.0, 1.25),
|
325 |
+
elem_classes="viewport",
|
326 |
+
label="3D preview (low-res, relief highlight)",
|
327 |
+
interactive=False,
|
328 |
+
)
|
329 |
+
files_3d = gr.Files(
|
330 |
+
label="3D model outputs (high-res)",
|
331 |
+
elem_id="download",
|
332 |
+
interactive=False,
|
333 |
+
)
|
334 |
+
|
335 |
+
blocks_settings_depth = [ensemble_size, denoise_steps, processing_res]
|
336 |
+
blocks_settings_3d = [plane_near, plane_far, embossing, size_longest_px, size_longest_cm, filter_size,
|
337 |
+
frame_thickness, frame_near, frame_far]
|
338 |
+
blocks_settings = blocks_settings_depth + blocks_settings_3d
|
339 |
+
map_id_to_default = {b._id: b.value for b in blocks_settings}
|
340 |
+
|
341 |
+
inputs = [
|
342 |
+
input_image,
|
343 |
+
ensemble_size,
|
344 |
+
denoise_steps,
|
345 |
+
processing_res,
|
346 |
+
input_output_16bit,
|
347 |
+
input_output_fp32,
|
348 |
+
input_output_vis,
|
349 |
+
plane_near,
|
350 |
+
plane_far,
|
351 |
+
embossing,
|
352 |
+
filter_size,
|
353 |
+
frame_near,
|
354 |
+
]
|
355 |
+
outputs = [
|
356 |
+
submit_btn,
|
357 |
+
input_image,
|
358 |
+
output_slider,
|
359 |
+
files,
|
360 |
+
]
|
361 |
+
|
362 |
+
def submit_depth_fn(*args):
|
363 |
+
out = list(process_pipe(*args))
|
364 |
+
out = [gr.Button(interactive=False), gr.Image(interactive=False)] + out
|
365 |
+
return out
|
366 |
+
|
367 |
+
submit_btn.click(
|
368 |
+
fn=submit_depth_fn,
|
369 |
+
inputs=inputs,
|
370 |
+
outputs=outputs,
|
371 |
+
concurrency_limit=1,
|
372 |
+
)
|
373 |
+
|
374 |
+
demo_3d_header.render()
|
375 |
+
demo_3d.render()
|
376 |
+
|
377 |
+
def clear_fn():
|
378 |
+
out = []
|
379 |
+
for b in blocks_settings:
|
380 |
+
out.append(map_id_to_default[b._id])
|
381 |
+
out += [
|
382 |
+
gr.Button(interactive=True),
|
383 |
+
gr.Button(interactive=True),
|
384 |
+
gr.Image(value=None, interactive=True),
|
385 |
+
None, None, None, None, None, None, None,
|
386 |
+
]
|
387 |
+
return out
|
388 |
+
|
389 |
+
clear_btn.click(
|
390 |
+
fn=clear_fn,
|
391 |
+
inputs=[],
|
392 |
+
outputs=blocks_settings + [
|
393 |
+
submit_btn,
|
394 |
+
submit_3d,
|
395 |
+
input_image,
|
396 |
+
input_output_16bit,
|
397 |
+
input_output_fp32,
|
398 |
+
input_output_vis,
|
399 |
+
output_slider,
|
400 |
+
files,
|
401 |
+
viewer_3d,
|
402 |
+
files_3d,
|
403 |
+
],
|
404 |
+
)
|
405 |
+
|
406 |
+
def submit_3d_fn(*args):
|
407 |
+
out = list(process_3d(*args))
|
408 |
+
out = [gr.Button(interactive=False)] + out
|
409 |
+
return out
|
410 |
+
|
411 |
+
submit_3d.click(
|
412 |
+
fn=submit_3d_fn,
|
413 |
+
inputs=[
|
414 |
+
input_image,
|
415 |
+
files,
|
416 |
+
size_longest_px,
|
417 |
+
size_longest_cm,
|
418 |
+
filter_size,
|
419 |
+
plane_near,
|
420 |
+
plane_far,
|
421 |
+
embossing,
|
422 |
+
frame_thickness,
|
423 |
+
frame_near,
|
424 |
+
frame_far,
|
425 |
+
],
|
426 |
+
outputs=[submit_3d, viewer_3d, files_3d],
|
427 |
+
concurrency_limit=1,
|
428 |
+
)
|
429 |
+
|
430 |
+
def clear_3d_fn():
|
431 |
+
return [gr.Button(interactive=True), None, None]
|
432 |
+
|
433 |
+
clear_3d.click(
|
434 |
+
fn=clear_3d_fn,
|
435 |
+
inputs=[],
|
436 |
+
outputs=[submit_3d, viewer_3d, files_3d],
|
437 |
+
)
|
438 |
+
|
439 |
+
demo.queue(
|
440 |
+
api_open=False,
|
441 |
+
).launch(
|
442 |
+
server_name="0.0.0.0",
|
443 |
+
server_port=7860,
|
444 |
+
)
|
445 |
+
|
446 |
+
|
447 |
+
|
448 |
+
|
449 |
+
def main():
|
450 |
+
|
451 |
+
if os.path.isdir(REPO_DIR):
|
452 |
+
shutil.rmtree(REPO_DIR)
|
453 |
+
repo = git.Repo.clone_from(REPO_URL, REPO_DIR)
|
454 |
+
repo.git.checkout(REPO_HASH)
|
455 |
+
|
456 |
+
sys.path.append(os.path.join(os.getcwd(), REPO_DIR))
|
457 |
+
|
458 |
+
from diffusers import DiffusionPipeline
|
459 |
+
pipeline = DiffusionPipeline.from_pretrained("JUGGHM/temp_repo")
|
460 |
+
|
461 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
462 |
+
|
463 |
+
pipe = pipeline.from_pretrained(CHECKPOINT)
|
464 |
+
try:
|
465 |
+
import xformers
|
466 |
+
pipe.enable_xformers_memory_efficient_attention()
|
467 |
+
except:
|
468 |
+
pass # run without xformers
|
469 |
+
|
470 |
+
pipe = pipe.to(device)
|
471 |
+
run_demo_server(pipe)
|
472 |
+
|
473 |
+
|
474 |
+
if __name__ == "__main__":
|
475 |
+
main()
|