import gradio as gr import torch from PIL import Image from depthmaster import DepthMasterPipeline from depthmaster.modules.unet_2d_condition import UNet2DConditionModel def load_example(example_image): # 返回选中的图片 return example_image device = "cuda" if torch.cuda.is_available() else "cpu" model_repo_id = "zysong212/DepthMaster" # Replace to the model you would like to use # if torch.cuda.is_available(): # torch_dtype = torch.float16 # else: torch_dtype = torch.float32 # pipe = DepthMasterPipeline.from_pretrained('eval', torch_dtype=torch_dtype) # unet = UNet2DConditionModel.from_pretrained(os.path.join('eval', f'unet')) pipe = DepthMasterPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype) unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder="unet", torch_dtype=torch_dtype) pipe.unet = unet try: pipe.enable_xformers_memory_efficient_attention() except ImportError: pass # run without xformers pipe = pipe.to(device) # MAX_SEED = np.iinfo(np.int32).max # MAX_IMAGE_SIZE = 1024 # @spaces.GPU #[uncomment to use ZeroGPU] def infer( input_image, progress=gr.Progress(track_tqdm=True), ): # if randomize_seed: # seed = random.randint(0, MAX_SEED) # generator = torch.Generator().manual_seed(seed) # image = pipe( # prompt=prompt, # negative_prompt=negative_prompt, # guidance_scale=guidance_scale, # num_inference_steps=num_inference_steps, # width=width, # height=height, # generator=generator, # ).images[0] pipe_out = pipe( input_image, processing_res=768, match_input_res=True, batch_size=1, color_map="Spectral", show_progress_bar=True, resample_method="bilinear", ) # depth_pred: np.ndarray = pipe_out.depth_np depth_colored: Image.Image = pipe_out.depth_colored return depth_colored # 默认图像路径 example_images = [ "wild_example/000000000776.jpg", "wild_example/800x.jpg", "wild_example/000000055950.jpg", "wild_example/53441037037_c2cbd91ad2_k.jpg", "wild_example/53501906161_6109e3da29_b.jpg", "wild_example/m_1e31af1c.jpg", "wild_example/sg-11134201-7rd5x-lvlh48byidbqca.jpg" ] css = """ #col-container { margin: 0 auto; max-width: 640px; } #example-gallery { height: 80px; /* 设置缩略图高度 */ width: auto; /* 保持宽高比 */ margin: 0 auto; /* 图片间距 */ cursor: pointer; /* 鼠标指针变为手型 */ } """ with gr.Blocks(css=css) as demo: gr.Markdown("# DepthMaster") gr.Markdown("Official demo for DepthMaster. Please refer to our [paper](https://arxiv.org/abs/2501.02576), [project page](https://indu1ge.github.io/DepthMaster_page/), and [github](https://github.com/indu1ge/DepthMaster) for more details.") gr.Markdown(" ### Depth Estimation with DepthMaster.") # with gr.Column(elem_id="col-container"): # gr.Markdown(" # Depth Estimation") with gr.Row(): with gr.Column(): input_image = gr.Image(label="Input Image", type="pil", elem_id="input-image", interactive=True) with gr.Column(): depth_map = gr.Image(label="Depth Map with Slider View", type="pil", interactive=False, elem_id="depth-map") # 计算按钮 compute_button = gr.Button("Compute Depth") # # 添加示例图片选择器 # with gr.Row(): # gr.Markdown("### example images") # with gr.Row(elem_id="example-gallery"): # example_gallery = gr.Gallery( # label="", # value=example_images, # elem_id="example-gallery", # show_label=False, # interactive=True, # columns=10 # ) # 设置默认图片点击后的操作 # example_gallery.select( # fn=lambda img_path: img_path, # 回调函数：返回选择的路径 # inputs=[], # outputs=input_image # 输出设置为 Input Image # ) # example_gallery.click( # fn=load_example, # 选择图片的回调 # inputs=[example_gallery], # 输入：用户点击的图片 # outputs=[input_image] # 输出：更新 Input Image # ) # 设置计算按钮的回调 compute_button.click( fn=infer, # 回调函数 inputs=input_image, # 输入 outputs=depth_map # 输出 ) # 启动 Gradio 应用 demo.launch() # with gr.Column(scale=45): # img_in = gr.Image(type="pil") # with gr.Column(scale=45): # img_out = # with gr.Row(): # prompt = gr.Text( # label="Prompt", # show_label=False, # max_lines=1, # placeholder="Enter your prompt", # container=False, # ) # run_button = gr.Button("Run", scale=0, variant="primary") # result = gr.Image(label="Result", show_label=False) # with gr.Accordion("Advanced Settings", open=False): # negative_prompt = gr.Text( # label="Negative prompt", # max_lines=1, # placeholder="Enter a negative prompt", # visible=False, # ) # seed = gr.Slider( # label="Seed", # minimum=0, # maximum=MAX_SEED, # step=1, # value=0, # ) # randomize_seed = gr.Checkbox(label="Randomize seed", value=True) # with gr.Row(): # width = gr.Slider( # label="Width", # minimum=256, # maximum=MAX_IMAGE_SIZE, # step=32, # value=1024, # Replace with defaults that work for your model # ) # height = gr.Slider( # label="Height", # minimum=256, # maximum=MAX_IMAGE_SIZE, # step=32, # value=1024, # Replace with defaults that work for your model # ) # with gr.Row(): # guidance_scale = gr.Slider( # label="Guidance scale", # minimum=0.0, # maximum=10.0, # step=0.1, # value=0.0, # Replace with defaults that work for your model # ) # num_inference_steps = gr.Slider( # label="Number of inference steps", # minimum=1, # maximum=50, # step=1, # value=2, # Replace with defaults that work for your model # ) # gr.Examples(examples=examples, inputs=[prompt]) # gr.on( # triggers=[run_button.click, prompt.submit], # fn=infer, # inputs=[ # prompt, # negative_prompt, # seed, # randomize_seed, # # width, # # height, # # guidance_scale, # # num_inference_steps, # ], # outputs=[result, seed], # ) # if __name__ == "__main__": # demo.launch()