Spaces:
Running
on
Zero
Running
on
Zero
| import functools | |
| import os | |
| import shutil | |
| import sys | |
| import git | |
| import gradio as gr | |
| import numpy as np | |
| import torch as torch | |
| from PIL import Image | |
| from gradio_imageslider import ImageSlider | |
| import spaces | |
| import fire | |
| REPO_URL = "https://github.com/lemonaddie/geowizard.git" | |
| CHECKPOINT = "lemonaddie/Geowizard" | |
| REPO_DIR = "geowizard" | |
| if os.path.isdir(REPO_DIR): | |
| shutil.rmtree(REPO_DIR) | |
| repo = git.Repo.clone_from(REPO_URL, REPO_DIR) | |
| sys.path.append(os.path.join(os.getcwd(), REPO_DIR)) | |
| from pipeline.depth_normal_pipeline_clip_cfg import DepthNormalEstimationPipeline | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| pipe = DepthNormalEstimationPipeline.from_pretrained(CHECKPOINT) | |
| try: | |
| import xformers | |
| pipe.enable_xformers_memory_efficient_attention() | |
| except: | |
| pass # run without xformers | |
| pipe = pipe.to(device) | |
| #run_demo_server(pipe) | |
| def depth_normal(img, | |
| denoising_steps, | |
| ensemble_size, | |
| processing_res, | |
| guidance_scale, | |
| domain): | |
| #img = img.resize((processing_res, processing_res), Image.Resampling.LANCZOS) | |
| pipe_out = pipe( | |
| img, | |
| denoising_steps=denoising_steps, | |
| ensemble_size=ensemble_size, | |
| processing_res=processing_res, | |
| batch_size=0, | |
| guidance_scale=guidance_scale, | |
| domain=domain, | |
| show_progress_bar=True, | |
| ) | |
| depth_colored = pipe_out.depth_colored | |
| normal_colored = pipe_out.normal_colored | |
| return depth_colored, normal_colored | |
| def run_demo(): | |
| custom_theme = gr.themes.Soft(primary_hue="blue").set( | |
| button_secondary_background_fill="*neutral_100", | |
| button_secondary_background_fill_hover="*neutral_200") | |
| custom_css = '''#disp_image { | |
| text-align: center; /* Horizontally center the content */ | |
| }''' | |
| _TITLE = '''GeoWizard: Unleashing the Diffusion Priors for 3D Geometry Estimation from a Single Image''' | |
| _DESCRIPTION = ''' | |
| <div> | |
| Generate consistent depth and normal from single image. High quality and rich details. | |
| <a style="display:inline-block; margin-left: .5em" href='https://github.com/fuxiao0719/GeoWizard/'><img src='https://img.shields.io/github/stars/fuxiao0719/GeoWizard?style=social' /></a> | |
| </div> | |
| ''' | |
| _GPU_ID = 0 | |
| with gr.Blocks(title=_TITLE, theme=custom_theme, css=custom_css) as demo: | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown('# ' + _TITLE) | |
| gr.Markdown(_DESCRIPTION) | |
| with gr.Row(variant='panel'): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(type='pil', image_mode='RGBA', height=320, label='Input image') | |
| example_folder = os.path.join(os.path.dirname(__file__), "./files") | |
| example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)] | |
| gr.Examples( | |
| examples=example_fns, | |
| inputs=[input_image], | |
| # outputs=[input_image], | |
| cache_examples=False, | |
| label='Examples (click one of the images below to start)', | |
| examples_per_page=30 | |
| ) | |
| with gr.Column(scale=1): | |
| with gr.Accordion('Advanced options', open=True): | |
| with gr.Column(): | |
| domain = gr.Radio( | |
| [ | |
| ("Outdoor", "outdoor"), | |
| ("Indoor", "indoor"), | |
| ("Object", "object"), | |
| ], | |
| label="Data Type (Must Select One matches your image)", | |
| value="indoor", | |
| ) | |
| guidance_scale = gr.Slider( | |
| label="Classifier Free Guidance Scale", | |
| minimum=1, | |
| maximum=5, | |
| step=1, | |
| value=3, | |
| ) | |
| denoising_steps = gr.Slider( | |
| label="Number of denoising steps (More stepes, better quality)", | |
| minimum=1, | |
| maximum=50, | |
| step=1, | |
| value=20, | |
| ) | |
| ensemble_size = gr.Slider( | |
| label="Ensemble size (1 will be enough. More steps, higher accuracy)", | |
| minimum=1, | |
| maximum=15, | |
| step=1, | |
| value=1, | |
| ) | |
| processing_res = gr.Radio( | |
| [ | |
| ("Native", 0), | |
| ("Recommended", 768), | |
| ], | |
| label="Processing resolution", | |
| value=768, | |
| ) | |
| run_btn = gr.Button('Generate', variant='primary', interactive=True) | |
| with gr.Row(): | |
| with gr.Column(): | |
| depth = gr.Image(interactive=False, show_label=False) | |
| with gr.Column(): | |
| normal = gr.Image(interactive=False, show_label=False) | |
| run_btn.click(fn=depth_normal, | |
| inputs=[input_image, denoising_steps, | |
| ensemble_size, | |
| processing_res, | |
| guidance_scale, | |
| domain], | |
| outputs=[depth, normal] | |
| ) | |
| demo.queue().launch(share=True, max_threads=80) | |
| if __name__ == '__main__': | |
| fire.Fire(run_demo) | |