diff --git a/README.md b/README.md index 9983225..6f4ee0d 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,7 @@ CUDA_VISIBLE_DEVICES=0 nohup python apps/image_to_3d.py > /dev/null 2>&1 & ### API Generate a 3D model from an image using the command-line API. - +Models will be downloaded automatically, please wait for the first run. ```sh python3 embodied_gen/scripts/imageto3d.py \ --image_path apps/assets/example_image/sample_04.jpg apps/assets/example_image/sample_19.jpg \ @@ -87,7 +87,7 @@ python apps/text_to_3d.py ``` ### API - +Models will be downloaded automatically, see `download_kolors_weights`. ```sh bash embodied_gen/scripts/textto3d.sh \ --prompts "small bronze figurine of a lion" "带木质底座，具有经纬线的地球仪" "橙色电动手钻，有磨损细节" \ @@ -109,8 +109,7 @@ python apps/texture_edit.py ``` ### API -Generate textures for a 3D mesh using a text prompt. - +Models will be downloaded automatically, see `download_kolors_weights`, `geo_cond_mv`. ```sh bash embodied_gen/scripts/texture_gen.sh \ --mesh_path "apps/assets/example_texture/meshes/robot_text.obj" \ diff --git a/apps/common.py b/apps/common.py index bbef969..9ecea89 100644 --- a/apps/common.py +++ b/apps/common.py @@ -136,32 +136,6 @@ def patched_setup_functions(self): Gaussian.setup_functions = patched_setup_functions -def download_kolors_weights() -> None: - logger.info(f"Download kolors weights from huggingface...") - subprocess.run( - [ - "huggingface-cli", - "download", - "--resume-download", - "Kwai-Kolors/Kolors", - "--local-dir", - "weights/Kolors", - ], - check=True, - ) - subprocess.run( - [ - "huggingface-cli", - "download", - "--resume-download", - "Kwai-Kolors/Kolors-IP-Adapter-Plus", - "--local-dir", - "weights/Kolors-IP-Adapter-Plus", - ], - check=True, - ) - - if os.getenv("GRADIO_APP") == "imageto3d": RBG_REMOVER = RembgRemover() RBG14_REMOVER = BMGG14Remover() @@ -185,9 +159,6 @@ elif os.getenv("GRADIO_APP") == "textto3d": ) # PIPELINE.cuda() text_model_dir = "weights/Kolors" - if not os.path.exists(text_model_dir): - download_kolors_weights() - PIPELINE_IMG_IP = build_text2img_ip_pipeline(text_model_dir, ref_scale=0.3) PIPELINE_IMG = build_text2img_pipeline(text_model_dir) SEG_CHECKER = ImageSegChecker(GPT_CLIENT) @@ -198,9 +169,6 @@ elif os.getenv("GRADIO_APP") == "textto3d": os.path.dirname(os.path.abspath(__file__)), "sessions/textto3d" ) elif os.getenv("GRADIO_APP") == "texture_edit": - if not os.path.exists("weights/Kolors"): - download_kolors_weights() - PIPELINE_IP = build_texture_gen_pipe( base_ckpt_dir="./weights", ip_adapt_scale=0.7, diff --git a/apps/image_to_3d.py b/apps/image_to_3d.py index a4d655a..f15e263 100644 --- a/apps/image_to_3d.py +++ b/apps/image_to_3d.py @@ -153,8 +153,8 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo: label="Randomize Seed", value=False ) project_delight = gr.Checkbox( - label="Backproject delighting", - value=False, + label="Back-project Delight", + value=True, ) gr.Markdown("Geo Structure Generation") with gr.Row(): diff --git a/apps/text_to_3d.py b/apps/text_to_3d.py index d0c2292..f13bfb8 100644 --- a/apps/text_to_3d.py +++ b/apps/text_to_3d.py @@ -152,7 +152,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo: label="Randomize Seed", value=False ) project_delight = gr.Checkbox( - label="backproject delight", value=False + label="Back-project Delight", value=True ) gr.Markdown("Geo Structure Generation") with gr.Row(): diff --git a/apps/texture_edit.py b/apps/texture_edit.py index e65a556..89e2095 100644 --- a/apps/texture_edit.py +++ b/apps/texture_edit.py @@ -215,7 +215,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo: ) with gr.Row(): project_delight = gr.Checkbox( - label="delight", value=True + label="Back-project delight", value=True ) fix_mesh = gr.Checkbox( label="simplify mesh", value=False diff --git a/embodied_gen/models/text_model.py b/embodied_gen/models/text_model.py index 16857de..6379b16 100644 --- a/embodied_gen/models/text_model.py +++ b/embodied_gen/models/text_model.py @@ -16,7 +16,9 @@ import logging +import os import random +import subprocess import numpy as np import torch @@ -47,14 +49,46 @@ __all__ = [ "build_text2img_ip_pipeline", "build_text2img_pipeline", "text2img_gen", + "download_kolors_weights", ] +def download_kolors_weights(local_dir: str = "weights/Kolors") -> None: + logger.info(f"Download kolors weights from huggingface...") + os.makedirs(local_dir, exist_ok=True) + subprocess.run( + [ + "huggingface-cli", + "download", + "--resume-download", + "Kwai-Kolors/Kolors", + "--local-dir", + local_dir, + ], + check=True, + ) + + ip_adapter_path = f"{local_dir}/../Kolors-IP-Adapter-Plus" + subprocess.run( + [ + "huggingface-cli", + "download", + "--resume-download", + "Kwai-Kolors/Kolors-IP-Adapter-Plus", + "--local-dir", + ip_adapter_path, + ], + check=True, + ) + + def build_text2img_ip_pipeline( ckpt_dir: str, ref_scale: float, device: str = "cuda", ) -> StableDiffusionXLPipelineIP: + download_kolors_weights(ckpt_dir) + text_encoder = ChatGLMModel.from_pretrained( f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16 ).half() @@ -106,6 +140,8 @@ def build_text2img_pipeline( ckpt_dir: str, device: str = "cuda", ) -> StableDiffusionXLPipeline: + download_kolors_weights(ckpt_dir) + text_encoder = ChatGLMModel.from_pretrained( f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16 ).half() diff --git a/embodied_gen/models/texture_model.py b/embodied_gen/models/texture_model.py index dd2b5d5..1745163 100644 --- a/embodied_gen/models/texture_model.py +++ b/embodied_gen/models/texture_model.py @@ -28,6 +28,7 @@ from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import ( StableDiffusionXLControlNetImg2ImgPipeline, ) from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection +from embodied_gen.models.text_model import download_kolors_weights __all__ = [ "build_texture_gen_pipe", @@ -40,6 +41,8 @@ def build_texture_gen_pipe( ip_adapt_scale: float = 0, device: str = "cuda", ) -> DiffusionPipeline: + download_kolors_weights(f"{base_ckpt_dir}/Kolors") + tokenizer = ChatGLMTokenizer.from_pretrained( f"{base_ckpt_dir}/Kolors/text_encoder" )