File size: 3,927 Bytes
6720d7b 74d5114 c5e4846 74d5114 a6267f6 |
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 |
import os
os.system("pip freeze")
import spaces
import gradio as gr
import torch as torch
from diffusers import MarigoldDepthPipeline, DDIMScheduler
from gradio_dualvision import DualVisionApp
from huggingface_hub import login
from PIL import Image
CHECKPOINT = "prs-eth/marigold-depth-v1-1"
if "Gty20030709" in os.environ:
login(token=os.environ["Gty20030709"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
pipe = MarigoldDepthPipeline.from_pretrained(CHECKPOINT)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
pipe = pipe.to(device=device, dtype=dtype)
try:
import xformers
pipe.enable_xformers_memory_efficient_attention()
except:
pass
class MarigoldDepthApp(DualVisionApp):
DEFAULT_SEED = 2024
DEFAULT_ENSEMBLE_SIZE = 1
DEFAULT_DENOISE_STEPS = 4
DEFAULT_PROCESSING_RES = 768
def make_header(self):
gr.Markdown(
"""
<h2><a href="https://huggingface.co/spaces/prs-eth/marigold" style="color: black;">Marigold Depth Estimation</a></h2>
"""
)
with gr.Row(elem_classes="remove-elements"):
gr.Markdown(
)
def build_user_components(self):
with gr.Column():
ensemble_size = gr.Slider(
label="Ensemble size",
minimum=1,
maximum=10,
step=1,
value=self.DEFAULT_ENSEMBLE_SIZE,
)
denoise_steps = gr.Slider(
label="Number of denoising steps",
minimum=1,
maximum=20,
step=1,
value=self.DEFAULT_DENOISE_STEPS,
)
processing_res = gr.Radio(
[
("Native", 0),
("Recommended", 768),
],
label="Processing resolution",
value=self.DEFAULT_PROCESSING_RES,
)
return {
"ensemble_size": ensemble_size,
"denoise_steps": denoise_steps,
"processing_res": processing_res,
}
def process(self, image_in: Image.Image, **kwargs):
ensemble_size = kwargs.get("ensemble_size", self.DEFAULT_ENSEMBLE_SIZE)
denoise_steps = kwargs.get("denoise_steps", self.DEFAULT_DENOISE_STEPS)
processing_res = kwargs.get("processing_res", self.DEFAULT_PROCESSING_RES)
generator = torch.Generator(device=device).manual_seed(self.DEFAULT_SEED)
pipe_out = pipe(
image_in,
ensemble_size=ensemble_size,
num_inference_steps=denoise_steps,
processing_resolution=processing_res,
batch_size=1 if processing_res == 0 else 2,
output_uncertainty=ensemble_size >= 3,
generator=generator,
)
depth_vis = pipe.image_processor.visualize_depth(pipe_out.prediction)[0]
depth_16bit = pipe.image_processor.export_depth_to_16bit_png(pipe_out.prediction)[0]
out_modalities = {
"Depth Visualization": depth_vis,
"Depth 16-bit": depth_16bit,
}
if ensemble_size >= 3:
uncertainty = pipe.image_processor.visualize_uncertainty(pipe_out.uncertainty)[0]
out_modalities["Uncertainty"] = uncertainty
out_settings = {
"ensemble_size": ensemble_size,
"denoise_steps": denoise_steps,
"processing_res": processing_res,
}
return out_modalities, out_settings
with MarigoldDepthApp(
title="Marigold Depth",
examples_path="files",
examples_per_page=12,
squeeze_canvas=True,
spaces_zero_gpu_enabled=True,
) as demo:
demo.queue(
api_open=False,
).launch(
server_name="0.0.0.0",
server_port=7860,
)
|