Upload folder using huggingface_hub

.ipynb_checkpoints/README-checkpoint.md

@@ -0,0 +1,231 @@
+<!-- ## **Hunyuan3D-1.0** -->
+
+<p align="center">
+  <img src="./assets/logo.png"  height=200>
+</p>
+
+# Tencent Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D Generation
+
+<div align="center">
+  <a href="https://github.com/tencent/Hunyuan3D-1"><img src="https://img.shields.io/static/v1?label=Code&message=Github&color=blue&logo=github-pages"></a> &ensp;
+  <a href="https://3d.hunyuan.tencent.com"><img src="https://img.shields.io/static/v1?label=Homepage&message=Tencent Hunyuan3D&color=blue&logo=github-pages"></a> &ensp;
+  <a href="https://arxiv.org/pdf/2411.02293"><img src="https://img.shields.io/static/v1?label=Tech Report&message=Arxiv&color=red&logo=arxiv"></a> &ensp;
+  <a href="https://huggingface.co/Tencent/Hunyuan3D-1"><img src="https://img.shields.io/static/v1?label=Checkpoints&message=HuggingFace&color=yellow"></a> &ensp;
+  <a href="https://huggingface.co/spaces/Tencent/Hunyuan3D-1"><img src="https://img.shields.io/static/v1?label=Demo&message=HuggingFace&color=yellow"></a> &ensp;
+</div>
+
+
+## 🔥🔥🔥 News!!
+
+* Nov 5, 2024: 💬 We support demo running  image_to_3d generation now. Please check the [script](#using-gradio) below.
+* Nov 5, 2024: 💬 We support demo running  text_to_3d generation now. Please check the [script](#using-gradio) below.
+
+
+## 📑 Open-source Plan
+
+- [x] Inference 
+- [x] Checkpoints
+- [ ] Baking related
+- [ ] Training
+- [ ] ComfyUI
+- [ ] Distillation Version
+- [ ] TensorRT Version
+
+
+
+## **Abstract**
+<p align="center">
+  <img src="./assets/teaser.png"  height=450>
+</p>
+
+While 3D generative models have greatly improved artists' workflows, the existing diffusion models for 3D generation suffer from slow generation and poor generalization. To address this issue, we propose a two-stage approach named Hunyuan3D-1.0 including a lite version and a standard version, that both support text- and image-conditioned generation.
+
+In the first stage, we employ a multi-view diffusion model that efficiently generates multi-view RGB in approximately 4 seconds. These multi-view images capture rich details of the 3D asset from different viewpoints, relaxing the tasks from single-view to multi-view reconstruction. In the second stage, we introduce a feed-forward reconstruction model that rapidly and faithfully reconstructs the 3D asset given the generated multi-view images in approximately 7 seconds. The reconstruction network learns to handle noises and in-consistency introduced by the multi-view diffusion and leverages the available information from the condition image to efficiently recover the 3D structure.
+
+Our framework involves the text-to-image model, i.e., Hunyuan-DiT, making it a unified framework to support both text- and image-conditioned 3D generation. Our standard version has 3x more parameters than our lite and other existing model. Our Hunyuan3D-1.0 achieves an impressive balance between speed and quality, significantly reducing generation time while maintaining the quality and diversity of the produced assets.
+
+
+## 🎉 **Hunyuan3D-1 Architecture**
+
+<p align="center">
+  <img src="./assets/overview_3.png"  height=400>
+</p>
+
+
+## 📈 Comparisons
+
+We have evaluated Hunyuan3D-1.0 with other open-source 3d-generation methods, our Hunyuan3D-1.0 received the highest user preference across 5 metrics. Details in the picture on the lower left.
+
+The lite model takes around 10 seconds to produce a 3D mesh from a single image on an NVIDIA A100 GPU, while the standard model takes roughly 25 seconds. The plot laid out in the lower right demonstrates that Hunyuan3D-1.0 achieves an optimal balance between quality and efficiency.
+
+<p align="center">
+  <img src="./assets/radar.png"  height=300>
+  <img src="./assets/runtime.png"  height=300>
+</p>
+
+## Get Started
+
+#### Begin by cloning the repository:
+
+```shell
+git clone https://github.com/tencent/Hunyuan3D-1
+cd Hunyuan3D-1
+```
+
+#### Installation Guide for Linux
+
+We provide an env_install.sh script file for setting up environment. 
+
+```
+# step 1, create conda env
+conda create -n hunyuan3d-1 python=3.9 or 3.10 or 3.11 or 3.12
+conda activate hunyuan3d-1
+
+# step 2. install torch realated package
+which pip # check pip corresponds to python
+
+# modify the cuda version according to your machine (recommended)
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+
+# step 3. install other packages
+bash env_install.sh
+```
+<details>
+<summary>💡Other tips for envrionment installation</summary>
+    
+Optionally, you can install xformers or flash_attn to acclerate computation:
+
+```
+pip install xformers --index-url https://download.pytorch.org/whl/cu121
+```
+```
+pip install flash_attn
+```
+
+Most environment errors are caused by a mismatch between machine and packages. You can try manually specifying the version, as shown in the following successful cases:
+```
+# python3.9
+pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118 
+```
+
+when install pytorch3d, the gcc version is preferably greater than 9, and the gpu driver should not be too old.
+
+</details>
+
+#### Download Pretrained Models
+
+The models are available at [https://huggingface.co/tencent/Hunyuan3D-1](https://huggingface.co/tencent/Hunyuan3D-1):
+
++ `Hunyuan3D-1/lite`, lite model for multi-view generation.
++ `Hunyuan3D-1/std`, standard model for multi-view generation.
++ `Hunyuan3D-1/svrm`, sparse-view reconstruction model.
+
+
+To download the model, first install the huggingface-cli. (Detailed instructions are available [here](https://huggingface.co/docs/huggingface_hub/guides/cli).)
+
+```shell
+python3 -m pip install "huggingface_hub[cli]"
+```
+
+Then download the model using the following commands:
+
+```shell
+mkdir weights
+huggingface-cli download tencent/Hunyuan3D-1 --local-dir ./weights
+
+mkdir weights/hunyuanDiT
+huggingface-cli download Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled --local-dir ./weights/hunyuanDiT
+```
+
+#### Inference 
+For text to 3d generation, we supports bilingual Chinese and English, you can use the following command to inference.
+```python
+python3 main.py \
+    --text_prompt "a lovely rabbit" \
+    --save_folder ./outputs/test/ \
+    --max_faces_num 90000 \
+    --do_texture_mapping \
+    --do_render
+```
+
+For image to 3d generation, you can use the following command to inference.
+```python
+python3 main.py \
+    --image_prompt "/path/to/your/image" \
+    --save_folder ./outputs/test/ \
+    --max_faces_num 90000 \
+    --do_texture_mapping \
+    --do_render
+```
+We list some more useful configurations for easy usage:
+
+|    Argument        |  Default  |                     Description                     |
+|:------------------:|:---------:|:---------------------------------------------------:|
+|`--text_prompt`  |   None    |The text prompt for 3D generation         |
+|`--image_prompt` |   None    |The image prompt for 3D generation         |
+|`--t2i_seed`     |    0      |The random seed for generating images        |
+|`--t2i_steps`    |    25     |The number of steps for sampling of text to image  |
+|`--gen_seed`     |    0      |The random seed for generating 3d generation        |
+|`--gen_steps`    |    50     |The number of steps for sampling of 3d generation  |
+|`--max_faces_numm` | 90000  |The limit number of faces of 3d mesh |
+|`--save_memory`   | False   |module will move to cpu automatically|
+|`--do_texture_mapping` |   False    |Change vertex shadding to texture shading  |
+|`--do_render`  |   False   |render gif   |
+
+
+We have also prepared scripts with different configurations for reference
+- Inference Std-pipeline requires 30GB VRAM (24G VRAM with --save_memory).
+- Inference Lite-pipeline requires 22GB VRAM (18G VRAM with --save_memory).
+- Note: --save_memory will increase inference time
+
+```bash
+bash scripts/text_to_3d_std.sh 
+bash scripts/text_to_3d_lite.sh 
+bash scripts/image_to_3d_std.sh 
+bash scripts/image_to_3d_lite.sh 
+```
+
+If your gpu memory is 16G, you can try to run modules in pipeline seperately:
+```bash
+bash scripts/text_to_3d_std_separately.sh 'a lovely rabbit' ./outputs/test # >= 16G
+bash scripts/text_to_3d_lite_separately.sh 'a lovely rabbit' ./outputs/test # >= 14G
+bash scripts/image_to_3d_std_separately.sh ./demos/example_000.png ./outputs/test  # >= 16G
+bash scripts/image_to_3d_lite_separately.sh ./demos/example_000.png ./outputs/test # >= 10G
+```
+
+#### Using Gradio
+
+We have prepared two versions of multi-view generation, std and lite.
+
+```shell
+# std 
+python3 app.py
+python3 app.py --save_memory
+
+# lite
+python3 app.py --use_lite
+python3 app.py --use_lite --save_memory
+```
+
+Then the demo can be accessed through http://0.0.0.0:8080. It should be noted that the 0.0.0.0 here needs to be X.X.X.X with your server IP.
+
+## Camera Parameters
+
+Output views are a fixed set of camera poses:
+
++ Azimuth (relative to input view): `+0, +60, +120, +180, +240, +300`.
+
+
+## Citation
+
+If you found this repository helpful, please cite our report:
+```bibtex
+@misc{yang2024tencent,
+    title={Tencent Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D Generation},
+    author={Xianghui Yang and Huiwen Shi and Bowen Zhang and Fan Yang and Jiacheng Wang and Hongxu Zhao and Xinhai Liu and Xinzhou Wang and Qingxiang Lin and Jiaao Yu and Lifu Wang and Zhuo Chen and Sicong Liu and Yuhong Liu and Yong Yang and Di Wang and Jie Jiang and Chunchao Guo},
+    year={2024},
+    eprint={2411.02293},
+    archivePrefix={arXiv},
+    primaryClass={cs.CV}
+}
+```

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,343 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os
+import warnings
+import argparse
+import gradio as gr
+from glob import glob
+import shutil
+import torch
+import numpy as np
+from PIL import Image
+from einops import rearrange
+
+from infer import seed_everything, save_gif
+from infer import Text2Image, Removebg, Image2Views, Views2Mesh, GifRenderer
+
+warnings.simplefilter('ignore', category=UserWarning)
+warnings.simplefilter('ignore', category=FutureWarning)
+warnings.simplefilter('ignore', category=DeprecationWarning)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--use_lite", default=False, action="store_true")
+parser.add_argument("--mv23d_cfg_path", default="./svrm/configs/svrm.yaml", type=str)
+parser.add_argument("--mv23d_ckt_path", default="weights/svrm/svrm.safetensors", type=str)
+parser.add_argument("--text2image_path", default="weights/hunyuanDiT", type=str)
+parser.add_argument("--save_memory", default=False, action="store_true")
+parser.add_argument("--device", default="cuda:0", type=str)
+args = parser.parse_args()
+
+################################################################
+# initial setting
+################################################################
+
+CONST_PORT = 8080
+CONST_MAX_QUEUE = 1
+CONST_SERVER = '0.0.0.0'
+
+CONST_HEADER = '''
+<h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'><b>Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D
+Generationr</b></a></h2>
+Code: <a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'>GitHub</a>. Techenical report: <a href='https://arxiv.org/abs/placeholder' target='_blank'>ArXiv</a>.
+
+❗️❗️❗️**Important Notes:**
+- By default, our demo can export a .obj mesh with vertex colors or a .glb mesh.
+- If you select "texture mapping," it will export a .obj mesh with a texture map or a .glb mesh.
+- If you select "render GIF," it will export a GIF image rendering of the .glb file.
+- If the result is unsatisfactory, please try a different seed value (Default: 0).
+'''
+
+CONST_CITATION = r"""
+If HunYuan3D-1 is helpful, please help to ⭐ the <a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'>Github Repo</a>. Thanks! [![GitHub Stars](https://img.shields.io/github/stars/tencent/Hunyuan3D-1?style=social)](https://github.com/tencent/Hunyuan3D-1)
+---
+📝 **Citation**
+If you find our work useful for your research or applications, please cite using this bibtex:
+```bibtex
+@misc{yang2024tencent,
+    title={Tencent Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D Generation},
+    author={Xianghui Yang and Huiwen Shi and Bowen Zhang and Fan Yang and Jiacheng Wang and Hongxu Zhao and Xinhai Liu and Xinzhou Wang and Qingxiang Lin and Jiaao Yu and Lifu Wang and Zhuo Chen and Sicong Liu and Yuhong Liu and Yong Yang and Di Wang and Jie Jiang and Chunchao Guo},
+    year={2024},
+    eprint={2411.02293},
+    archivePrefix={arXiv},
+    primaryClass={cs.CV}
+}
+```
+"""
+
+################################################################
+# prepare text examples and image examples
+################################################################
+
+def get_example_img_list():
+    print('Loading example img list ...')
+    return sorted(glob('./demos/example_*.png'))
+
+def get_example_txt_list():
+    print('Loading example txt list ...')
+    txt_list  = list()
+    for line in open('./demos/example_list.txt'):
+        txt_list.append(line.strip())
+    return txt_list
+
+example_is = get_example_img_list()
+example_ts = get_example_txt_list()
+
+################################################################
+# initial models
+################################################################
+
+worker_xbg = Removebg()
+print(f"loading {args.text2image_path}")
+worker_t2i = Text2Image(
+    pretrain = args.text2image_path, 
+    device = args.device, 
+    save_memory = args.save_memory
+)
+worker_i2v = Image2Views(
+    use_lite = args.use_lite, 
+    device = args.device,
+    save_memory = args.save_memory
+)
+worker_v23 = Views2Mesh(
+    args.mv23d_cfg_path, 
+    args.mv23d_ckt_path, 
+    use_lite = args.use_lite, 
+    device = args.device,
+    save_memory = args.save_memory
+)
+worker_gif = GifRenderer(args.device)
+
+def stage_0_t2i(text, image, seed, step):
+    os.makedirs('./outputs/app_output', exist_ok=True)
+    exists = set(int(_) for _ in os.listdir('./outputs/app_output') if not _.startswith("."))
+    if len(exists) == 30: shutil.rmtree(f"./outputs/app_output/0");cur_id = 0
+    else:                 cur_id = min(set(range(30)) - exists)
+    if os.path.exists(f"./outputs/app_output/{(cur_id + 1) % 30}"):
+        shutil.rmtree(f"./outputs/app_output/{(cur_id + 1) % 30}")
+    save_folder = f'./outputs/app_output/{cur_id}'
+    os.makedirs(save_folder, exist_ok=True)
+
+    dst = save_folder + '/img.png'
+    
+    if not text:
+        if image is None: 
+            return dst, save_folder
+            raise gr.Error("Upload image or provide text ...")
+        image.save(dst)
+        return dst, save_folder
+        
+    image = worker_t2i(text, seed, step)
+    image.save(dst)
+    dst = worker_xbg(image, save_folder)
+    return dst, save_folder
+
+def stage_1_xbg(image, save_folder): 
+    if isinstance(image, str):
+        image = Image.open(image)
+    dst =  save_folder + '/img_nobg.png'
+    rgba = worker_xbg(image)
+    rgba.save(dst)
+    return dst
+    
+def stage_2_i2v(image, seed, step, save_folder):
+    if isinstance(image, str):
+        image = Image.open(image)
+    gif_dst = save_folder + '/views.gif'
+    res_img, pils = worker_i2v(image, seed, step)
+    save_gif(pils, gif_dst)
+    views_img, cond_img = res_img[0], res_img[1]
+    img_array = np.asarray(views_img, dtype=np.uint8)
+    show_img = rearrange(img_array, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
+    show_img = show_img[worker_i2v.order, ...]
+    show_img = rearrange(show_img, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
+    show_img = Image.fromarray(show_img) 
+    return views_img, cond_img, show_img
+
+def stage_3_v23(
+    views_pil, 
+    cond_pil, 
+    seed, 
+    save_folder,
+    target_face_count = 30000,
+    do_texture_mapping = True,
+    do_render =True
+): 
+    do_texture_mapping = do_texture_mapping or do_render
+    obj_dst = save_folder + '/mesh_with_colors.obj'
+    glb_dst = save_folder + '/mesh.glb'
+    worker_v23(
+        views_pil, 
+        cond_pil, 
+        seed = seed, 
+        save_folder = save_folder,
+        target_face_count = target_face_count,
+        do_texture_mapping = do_texture_mapping
+    )
+    return obj_dst, glb_dst
+
+def stage_4_gif(obj_dst, save_folder, do_render_gif=True):
+    if not do_render_gif: return None
+    gif_dst = save_folder + '/output.gif'
+    worker_gif(
+        save_folder + '/mesh.obj',
+        gif_dst_path = gif_dst
+    )
+    return gif_dst
+# ===============================================================
+# gradio display
+# ===============================================================
+with gr.Blocks() as demo:
+    gr.Markdown(CONST_HEADER)
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=2):
+            with gr.Tab("Text to 3D"):
+                with gr.Column():
+                    text = gr.TextArea('一只黑白相间的熊猫在白色背景上居中坐着，呈现出卡通风格和可爱氛围。', lines=1, max_lines=10, label='Input text')
+                    with gr.Row():
+                        textgen_seed = gr.Number(value=0, label="T2I seed", precision=0)
+                        textgen_step = gr.Number(value=25, label="T2I step", precision=0)
+                        textgen_SEED = gr.Number(value=0, label="Gen seed", precision=0)
+                        textgen_STEP = gr.Number(value=50, label="Gen step", precision=0)
+                        textgen_max_faces = gr.Number(value=90000, label="max number of faces", precision=0)
+                        
+                    with gr.Row():
+                        textgen_do_texture_mapping = gr.Checkbox(label="texture mapping", value=False, interactive=True)
+                        textgen_do_render_gif = gr.Checkbox(label="Render gif", value=False, interactive=True)
+                        textgen_submit = gr.Button("Generate", variant="primary")
+
+                    with gr.Row():
+                        gr.Examples(examples=example_ts, inputs=[text], label="Txt examples", examples_per_page=10)
+                    
+            with gr.Tab("Image to 3D"):
+                with gr.Column():
+                    input_image = gr.Image(label="Input image",
+                                           width=256, height=256, type="pil",
+                                           image_mode="RGBA", sources="upload",
+                                           interactive=True)
+                    with gr.Row(): 
+                        imggen_SEED = gr.Number(value=0, label="Gen seed", precision=0)
+                        imggen_STEP = gr.Number(value=50, label="Gen step", precision=0)
+                        imggen_max_faces = gr.Number(value=90000, label="max number of faces", precision=0)
+
+                    with gr.Row():
+                        imggen_do_texture_mapping = gr.Checkbox(label="texture mapping", value=False, interactive=True)
+                        imggen_do_render_gif = gr.Checkbox(label="Render gif", value=False, interactive=True)
+                        imggen_submit = gr.Button("Generate", variant="primary")       
+                    with gr.Row():
+                        gr.Examples(
+                            examples=example_is, 
+                            inputs=[input_image], 
+                            label="Img examples",
+                            examples_per_page=10
+                        )
+           
+        with gr.Column(scale=3):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    rem_bg_image = gr.Image(label="No backgraound image", type="pil",
+                                           image_mode="RGBA", interactive=False)
+                with gr.Column(scale=3):
+                    result_image = gr.Image(label="Multi views", type="pil", interactive=False)
+                
+            with gr.Row():                
+                result_3dobj = gr.Model3D(
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    label="Output Obj",
+                    show_label=True,
+                    visible=True,
+                    camera_position=[90, 90, None],
+                    interactive=False
+                )
+
+                result_3dglb = gr.Model3D(
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    label="Output Glb",
+                    show_label=True,
+                    visible=True,
+                    camera_position=[90, 90, None],
+                    interactive=False
+                )
+                result_gif = gr.Image(label="Rendered GIF", interactive=False)
+                
+            with gr.Row():    
+                gr.Markdown("""
+                We recommend downloading and opening Glb with 3D software, such as Blender, MeshLab, etc.
+                
+                Limited by gradio, Obj file here only be shown as vertex shading, but Glb can be texture shading.
+                """)
+
+#===============================================================
+# gradio running code
+#===============================================================
+
+    none = gr.State(None)
+    save_folder = gr.State()
+    cond_image = gr.State()
+    views_image = gr.State()
+    text_image = gr.State()
+    
+    textgen_submit.click(
+        fn=stage_0_t2i, inputs=[text, none, textgen_seed, textgen_step], 
+        outputs=[rem_bg_image, save_folder],
+    ).success(
+        fn=stage_2_i2v, inputs=[rem_bg_image, textgen_SEED, textgen_STEP, save_folder], 
+        outputs=[views_image, cond_image, result_image],
+    ).success(
+        fn=stage_3_v23, inputs=[views_image, cond_image, textgen_SEED, save_folder, 
+                                textgen_max_faces, textgen_do_texture_mapping,
+                                textgen_do_render_gif], 
+        outputs=[result_3dobj, result_3dglb],
+    ).success(
+        fn=stage_4_gif, inputs=[result_3dglb, save_folder, textgen_do_render_gif], 
+        outputs=[result_gif],
+    ).success(lambda: print('Text_to_3D Done ...'))
+
+    imggen_submit.click(
+        fn=stage_0_t2i, inputs=[none, input_image, textgen_seed, textgen_step], 
+        outputs=[text_image, save_folder],
+    ).success(
+        fn=stage_1_xbg, inputs=[text_image, save_folder], 
+        outputs=[rem_bg_image],
+    ).success(
+        fn=stage_2_i2v, inputs=[rem_bg_image, imggen_SEED, imggen_STEP, save_folder], 
+        outputs=[views_image, cond_image, result_image],
+    ).success(
+        fn=stage_3_v23, inputs=[views_image, cond_image, imggen_SEED, save_folder, 
+                                imggen_max_faces, imggen_do_texture_mapping, 
+                                imggen_do_render_gif], 
+        outputs=[result_3dobj, result_3dglb],
+    ).success(
+        fn=stage_4_gif, inputs=[result_3dglb, save_folder, imggen_do_render_gif], 
+        outputs=[result_gif],
+    ).success(lambda: print('Image_to_3D Done ...'))
+    
+#===============================================================
+# start gradio server
+#===============================================================
+
+    gr.Markdown(CONST_CITATION)
+    demo.queue(max_size=CONST_MAX_QUEUE)
+    demo.launch(server_name=CONST_SERVER, server_port=CONST_PORT)
+

.ipynb_checkpoints/app_hg-checkpoint.py

@@ -0,0 +1,384 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+import spaces
+import os
+os.environ['CUDA_HOME'] = '/usr/local/cuda-11*'
+import warnings
+import argparse
+import gradio as gr
+from glob import glob
+import shutil
+import torch
+import numpy as np
+from PIL import Image
+from einops import rearrange
+from huggingface_hub import snapshot_download
+
+from infer import seed_everything, save_gif
+from infer import Text2Image, Removebg, Image2Views, Views2Mesh, GifRenderer
+
+warnings.simplefilter('ignore', category=UserWarning)
+warnings.simplefilter('ignore', category=FutureWarning)
+warnings.simplefilter('ignore', category=DeprecationWarning)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--use_lite", default=False, action="store_true")
+parser.add_argument("--mv23d_cfg_path", default="./svrm/configs/svrm.yaml", type=str)
+parser.add_argument("--mv23d_ckt_path", default="weights/svrm/svrm.safetensors", type=str)
+parser.add_argument("--text2image_path", default="weights/hunyuanDiT", type=str)
+parser.add_argument("--save_memory", default=True) # , action="store_true")
+parser.add_argument("--device", default="cuda:0", type=str)
+args = parser.parse_args()
+
+def find_cuda():
+    # Check if CUDA_HOME or CUDA_PATH environment variables are set
+    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+
+    if cuda_home and os.path.exists(cuda_home):
+        return cuda_home
+
+    # Search for the nvcc executable in the system's PATH
+    nvcc_path = shutil.which('nvcc')
+
+    if nvcc_path:
+        # Remove the 'bin/nvcc' part to get the CUDA installation path
+        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
+        return cuda_path
+
+    return None
+
+cuda_path = find_cuda()
+
+if cuda_path:
+    print(f"CUDA installation found at: {cuda_path}")
+else:
+    print("CUDA installation not found")
+
+
+
+def download_models():
+    # Create weights directory if it doesn't exist
+    os.makedirs("weights", exist_ok=True)
+    os.makedirs("weights/hunyuanDiT", exist_ok=True)
+
+    # Download Hunyuan3D-1 model
+    try:
+        snapshot_download(
+            repo_id="tencent/Hunyuan3D-1",
+            local_dir="./weights",
+            resume_download=True
+        )
+        print("Successfully downloaded Hunyuan3D-1 model")
+    except Exception as e:
+        print(f"Error downloading Hunyuan3D-1: {e}")
+
+    # Download HunyuanDiT model
+    try:
+        snapshot_download(
+            repo_id="Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled",
+            local_dir="./weights/hunyuanDiT",
+            resume_download=True
+        )
+        print("Successfully downloaded HunyuanDiT model")
+    except Exception as e:
+        print(f"Error downloading HunyuanDiT: {e}")
+
+# Download models before starting the app
+download_models()
+
+################################################################
+
+CONST_PORT = 8080
+CONST_MAX_QUEUE = 1
+CONST_SERVER = '0.0.0.0'
+
+CONST_HEADER = '''
+<h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'><b>Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D
+Generationr</b></a></h2>
+Code: <a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'>GitHub</a>. Techenical report: <a href='https://arxiv.org/abs/placeholder' target='_blank'>ArXiv</a>.
+
+❗️❗️❗️**Important Notes:**
+- By default, our demo can export a .obj mesh with vertex colors or a .glb mesh.
+- If you select "texture mapping," it will export a .obj mesh with a texture map or a .glb mesh.
+- If you select "render GIF," it will export a GIF image rendering of the .glb file.
+- If the result is unsatisfactory, please try a different seed value (Default: 0).
+'''
+
+CONST_CITATION = r"""
+If HunYuan3D-1 is helpful, please help to ⭐ the <a href='https://github.com/tencent/Hunyuan3D-1' target='_blank'>Github Repo</a>. Thanks! [![GitHub Stars](https://img.shields.io/github/stars/tencent/Hunyuan3D-1?style=social)](https://github.com/tencent/Hunyuan3D-1)
+---
+📝 **Citation**
+If you find our work useful for your research or applications, please cite using this bibtex:
+```bibtex
+@misc{yang2024tencent,
+    title={Tencent Hunyuan3D-1.0: A Unified Framework for Text-to-3D and Image-to-3D Generation},
+    author={Xianghui Yang and Huiwen Shi and Bowen Zhang and Fan Yang and Jiacheng Wang and Hongxu Zhao and Xinhai Liu and Xinzhou Wang and Qingxiang Lin and Jiaao Yu and Lifu Wang and Zhuo Chen and Sicong Liu and Yuhong Liu and Yong Yang and Di Wang and Jie Jiang and Chunchao Guo},
+    year={2024},
+    eprint={2411.02293},
+    archivePrefix={arXiv},
+    primaryClass={cs.CV}
+}
+```
+"""
+
+################################################################
+
+def get_example_img_list():
+    print('Loading example img list ...')
+    return sorted(glob('./demos/example_*.png'))
+
+def get_example_txt_list():
+    print('Loading example txt list ...')
+    txt_list  = list()
+    for line in open('./demos/example_list.txt'):
+        txt_list.append(line.strip())
+    return txt_list
+
+example_is = get_example_img_list()
+example_ts = get_example_txt_list()
+################################################################
+
+worker_xbg = Removebg()
+print(f"loading {args.text2image_path}")
+worker_t2i = Text2Image(
+    pretrain = args.text2image_path, 
+    device = args.device, 
+    save_memory = args.save_memory
+)
+worker_i2v = Image2Views(
+    use_lite = args.use_lite, 
+    device = args.device,
+    save_memory = args.save_memory
+)
+worker_v23 = Views2Mesh(
+    args.mv23d_cfg_path, 
+    args.mv23d_ckt_path, 
+    use_lite = args.use_lite, 
+    device = args.device,
+    save_memory = args.save_memory
+)
+worker_gif = GifRenderer(args.device)
+
+@spaces.GPU
+def stage_0_t2i(text, image, seed, step):
+    os.makedirs('./outputs/app_output', exist_ok=True)
+    exists = set(int(_) for _ in os.listdir('./outputs/app_output') if not _.startswith("."))
+    if len(exists) == 30: shutil.rmtree(f"./outputs/app_output/0");cur_id = 0
+    else:                 cur_id = min(set(range(30)) - exists)
+    if os.path.exists(f"./outputs/app_output/{(cur_id + 1) % 30}"):
+        shutil.rmtree(f"./outputs/app_output/{(cur_id + 1) % 30}")
+    save_folder = f'./outputs/app_output/{cur_id}'
+    os.makedirs(save_folder, exist_ok=True)
+
+    dst = os.path.join(save_folder, 'img.png')
+    
+    if not text:
+        if image is None: 
+            return dst, save_folder
+            raise gr.Error("Upload image or provide text ...")
+        image.save(dst)
+        return dst, save_folder
+        
+    image = worker_t2i(text, seed, step)
+    image.save(dst)
+    dst = worker_xbg(image, save_folder)
+    return dst, save_folder
+    
+@spaces.GPU
+def stage_1_xbg(image, save_folder): 
+    if isinstance(image, str):
+        image = Image.open(image)
+    dst =  save_folder + '/img_nobg.png'
+    rgba = worker_xbg(image)
+    rgba.save(dst)
+    return dst
+    
+@spaces.GPU
+def stage_2_i2v(image, seed, step, save_folder):
+    if isinstance(image, str):
+        image = Image.open(image)
+    gif_dst = save_folder + '/views.gif'
+    res_img, pils = worker_i2v(image, seed, step)
+    save_gif(pils, gif_dst)
+    views_img, cond_img = res_img[0], res_img[1]
+    img_array = np.asarray(views_img, dtype=np.uint8)
+    show_img = rearrange(img_array, '(n h) (m w) c -> (n m) h w c', n=3, m=2)
+    show_img = show_img[worker_i2v.order, ...]
+    show_img = rearrange(show_img, '(n m) h w c -> (n h) (m w) c', n=2, m=3)
+    show_img = Image.fromarray(show_img) 
+    return views_img, cond_img, show_img
+
+@spaces.GPU
+def stage_3_v23(
+    views_pil, 
+    cond_pil, 
+    seed, 
+    save_folder,
+    target_face_count = 30000,
+    do_texture_mapping = True,
+    do_render =True
+): 
+    do_texture_mapping = do_texture_mapping or do_render
+    obj_dst = save_folder + '/mesh_with_colors.obj'
+    glb_dst = save_folder + '/mesh.glb'
+    worker_v23(
+        views_pil, 
+        cond_pil, 
+        seed = seed, 
+        save_folder = save_folder,
+        target_face_count = target_face_count,
+        do_texture_mapping = do_texture_mapping
+    )
+    return obj_dst, glb_dst
+
+@spaces.GPU
+def stage_4_gif(obj_dst, save_folder, do_render_gif=True):
+    if not do_render_gif: return None
+    gif_dst = save_folder + '/output.gif'
+    worker_gif(
+        save_folder + '/mesh.obj',
+        gif_dst_path = gif_dst
+    )
+    return gif_dst
+
+#===============================================================
+with gr.Blocks() as demo:
+    gr.Markdown(CONST_HEADER)
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=2):
+            with gr.Tab("Text to 3D"):
+                with gr.Column():
+                    text = gr.TextArea('一只黑白相间的熊猫在白色背景上居中坐着，呈现出卡通风格和可爱氛围。', lines=1, max_lines=10, label='Input text')
+                    with gr.Row():
+                        textgen_seed = gr.Number(value=0, label="T2I seed", precision=0)
+                        textgen_step = gr.Number(value=25, label="T2I step", precision=0)
+                        textgen_SEED = gr.Number(value=0, label="Gen seed", precision=0)
+                        textgen_STEP = gr.Number(value=50, label="Gen step", precision=0)
+                        textgen_max_faces = gr.Number(value=90000, label="max number of faces", precision=0)
+                        
+                    with gr.Row():
+                        textgen_do_texture_mapping = gr.Checkbox(label="texture mapping", value=False, interactive=True)
+                        textgen_do_render_gif = gr.Checkbox(label="Render gif", value=False, interactive=True)
+                        textgen_submit = gr.Button("Generate", variant="primary")
+
+                    with gr.Row():
+                        gr.Examples(examples=example_ts, inputs=[text], label="Txt examples", examples_per_page=10)
+                    
+            with gr.Tab("Image to 3D"):
+                with gr.Column():
+                    input_image = gr.Image(label="Input image",
+                                           width=256, height=256, type="pil",
+                                           image_mode="RGBA", sources="upload",
+                                           interactive=True)
+                    with gr.Row(): 
+                        imggen_SEED = gr.Number(value=0, label="Gen seed", precision=0)
+                        imggen_STEP = gr.Number(value=50, label="Gen step", precision=0)
+                        imggen_max_faces = gr.Number(value=90000, label="max number of faces", precision=0)
+
+                    with gr.Row():
+                        imggen_do_texture_mapping = gr.Checkbox(label="texture mapping", value=False, interactive=True)
+                        imggen_do_render_gif = gr.Checkbox(label="Render gif", value=False, interactive=True)
+                        imggen_submit = gr.Button("Generate", variant="primary")       
+                    with gr.Row():
+                        gr.Examples(examples=example_is, inputs=[input_image], label="Img examples", examples_per_page=10)
+           
+        with gr.Column(scale=3):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    rem_bg_image = gr.Image(label="No backgraound image", type="pil",
+                                           image_mode="RGBA", interactive=False)
+                with gr.Column(scale=3):
+                    result_image = gr.Image(label="Multi views", type="pil", interactive=False)
+                
+            with gr.Row():                
+                result_3dobj = gr.Model3D(
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    label="Output Obj",
+                    show_label=True,
+                    visible=True,
+                    camera_position=[90, 90, None],
+                    interactive=False
+                )
+
+                result_3dglb = gr.Model3D(
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    label="Output Glb",
+                    show_label=True,
+                    visible=True,
+                    camera_position=[90, 90, None],
+                    interactive=False
+                )
+                result_gif = gr.Image(label="Rendered GIF", interactive=False)
+                
+            with gr.Row():    
+                gr.Markdown("""
+                We recommend download and open Glb using 3D software, such as Blender, MeshLab, etc.
+                Limited by gradio, Obj file here only be shown as vertex shading, but Glb can be texture shading.
+                """)
+
+#===============================================================
+
+    none = gr.State(None)
+    save_folder = gr.State()
+    cond_image = gr.State()
+    views_image = gr.State()
+    text_image = gr.State()
+    
+    textgen_submit.click(
+        fn=stage_0_t2i, inputs=[text, none, textgen_seed, textgen_step], 
+        outputs=[rem_bg_image, save_folder],
+    ).success(
+        fn=stage_2_i2v, inputs=[rem_bg_image, textgen_SEED, textgen_STEP, save_folder], 
+        outputs=[views_image, cond_image, result_image],
+    ).success(
+        fn=stage_3_v23, inputs=[views_image, cond_image, textgen_SEED, save_folder, textgen_max_faces, textgen_do_texture_mapping, textgen_do_render_gif], 
+        outputs=[result_3dobj, result_3dglb],
+    ).success(
+        fn=stage_4_gif, inputs=[result_3dglb, save_folder, textgen_do_render_gif], 
+        outputs=[result_gif],
+    ).success(lambda: print('Text_to_3D Done ...'))
+
+    imggen_submit.click(
+        fn=stage_0_t2i, inputs=[none, input_image, textgen_seed, textgen_step], 
+        outputs=[text_image, save_folder],
+    ).success(
+        fn=stage_1_xbg, inputs=[text_image, save_folder], 
+        outputs=[rem_bg_image],
+    ).success(
+        fn=stage_2_i2v, inputs=[rem_bg_image, imggen_SEED, imggen_STEP, save_folder], 
+        outputs=[views_image, cond_image, result_image],
+    ).success(
+        fn=stage_3_v23, inputs=[views_image, cond_image, imggen_SEED, save_folder, imggen_max_faces, imggen_do_texture_mapping, imggen_do_render_gif], 
+        outputs=[result_3dobj, result_3dglb],
+    ).success(
+        fn=stage_4_gif, inputs=[result_3dglb, save_folder, imggen_do_render_gif], 
+        outputs=[result_gif],
+    ).success(lambda: print('Image_to_3D Done ...'))
+    
+#===============================================================
+
+    gr.Markdown(CONST_CITATION)
+    demo.queue(max_size=CONST_MAX_QUEUE)
+    demo.launch()
+

.ipynb_checkpoints/main-checkpoint.py

@@ -0,0 +1,164 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.l
+
+import os
+import warnings
+import torch
+from PIL import Image
+import argparse
+
+from infer import Text2Image, Removebg, Image2Views, Views2Mesh, GifRenderer
+
+warnings.simplefilter('ignore', category=UserWarning)
+warnings.simplefilter('ignore', category=FutureWarning)
+warnings.simplefilter('ignore', category=DeprecationWarning)
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--use_lite", default=False, action="store_true"
+    )
+    parser.add_argument(
+        "--mv23d_cfg_path", default="./svrm/configs/svrm.yaml", type=str
+    )
+    parser.add_argument(
+        "--mv23d_ckt_path", default="weights/svrm/svrm.safetensors", type=str
+    )
+    parser.add_argument(
+        "--text2image_path", default="weights/hunyuanDiT", type=str
+    )
+    parser.add_argument(
+        "--save_folder", default="./outputs/test/", type=str
+    )
+    parser.add_argument(
+        "--text_prompt", default="", type=str,
+    )
+    parser.add_argument(
+        "--image_prompt", default="", type=str
+    )
+    parser.add_argument(
+        "--device", default="cuda:0", type=str
+    )
+    parser.add_argument(
+        "--t2i_seed", default=0, type=int
+    )
+    parser.add_argument(
+        "--t2i_steps", default=25, type=int
+    )
+    parser.add_argument(
+        "--gen_seed", default=0, type=int
+    )
+    parser.add_argument(
+        "--gen_steps", default=50, type=int
+    )
+    parser.add_argument(
+        "--max_faces_num", default=80000, type=int, 
+        help="max num of face, suggest 80000 for effect, 10000 for speed"
+    )
+    parser.add_argument(
+        "--save_memory", default=False, action="store_true"
+    )
+    parser.add_argument(
+        "--do_texture_mapping", default=False, action="store_true"
+    )
+    parser.add_argument(
+        "--do_render", default=False, action="store_true"
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    
+    assert not (args.text_prompt and args.image_prompt), "Text and image can only be given to one"
+    assert args.text_prompt or args.image_prompt,        "Text and image can only be given to one"
+
+    # init model
+    rembg_model = Removebg()
+    image_to_views_model = Image2Views(
+        device=args.device, 
+        use_lite=args.use_lite,
+        save_memory=args.save_memory
+    )
+    
+    views_to_mesh_model = Views2Mesh(
+        args.mv23d_cfg_path, 
+        args.mv23d_ckt_path, 
+        args.device, 
+        use_lite=args.use_lite,
+        save_memory=args.save_memory
+    )
+    
+    if args.text_prompt:
+        text_to_image_model = Text2Image(
+            pretrain = args.text2image_path,
+            device = args.device, 
+            save_memory = args.save_memory
+        )
+    if args.do_render:
+        gif_renderer = GifRenderer(device=args.device)
+
+    # ---- ----- ---- ---- ---- ----
+
+    os.makedirs(args.save_folder, exist_ok=True)
+
+    # stage 1, text to image
+    if args.text_prompt:
+        res_rgb_pil = text_to_image_model(
+            args.text_prompt, 
+            seed=args.t2i_seed,  
+            steps=args.t2i_steps
+        )
+        res_rgb_pil.save(os.path.join(args.save_folder, "img.jpg"))
+    elif args.image_prompt:
+        res_rgb_pil = Image.open(args.image_prompt)
+
+    # stage 2, remove back ground
+    res_rgba_pil = rembg_model(res_rgb_pil)
+    res_rgb_pil.save(os.path.join(args.save_folder, "img_nobg.png"))
+
+    # stage 3, image to views
+    (views_grid_pil, cond_img), view_pil_list = image_to_views_model(
+        res_rgba_pil,
+        seed = args.gen_seed,
+        steps = args.gen_steps
+    )
+    views_grid_pil.save(os.path.join(args.save_folder, "views.jpg"))
+
+    # stage 4, views to mesh
+    views_to_mesh_model(
+        views_grid_pil, 
+        cond_img, 
+        seed = args.gen_seed,
+        target_face_count = args.max_faces_num,
+        save_folder = args.save_folder,
+        do_texture_mapping = args.do_texture_mapping
+    )
+
+    #  stage 5, render gif
+    if args.do_render:
+        gif_renderer(
+            os.path.join(args.save_folder, 'mesh.obj'),
+            gif_dst_path = os.path.join(args.save_folder, 'output.gif'),
+        )

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -0,0 +1,24 @@
+--find-links https://download.pytorch.org/whl/cu118
+torch==2.2.0
+torchvision==0.17.0
+diffusers
+numpy==1.26.4
+transformers
+rembg
+tqdm
+omegaconf
+matplotlib
+opencv-python
+imageio
+jaxtyping
+einops
+SentencePiece
+accelerate
+trimesh
+PyMCubes
+xatlas
+libigl
+git+https://github.com/facebookresearch/pytorch3d@stable
+git+https://github.com/NVlabs/nvdiffrast
+open3d
+ninja

README.md

@@ -1,14 +1,3 @@
----
-title: Hunyuan3D-1.0
-emoji: 😻
-colorFrom: purple
-colorTo: red
-sdk: gradio
-sdk_version: 5.5.0
-app_file: app_hg.py
-pinned: false
-short_description: Text-to-3D and Image-to-3D Generation
----
 <!-- ## **Hunyuan3D-1.0** -->
 
 <p align="center">

app.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 
@@ -47,6 +49,8 @@ parser.add_argument("--save_memory", default=False, action="store_true")
 parser.add_argument("--device", default="cuda:0", type=str)
 args = parser.parse_args()
 
+################################################################
+# initial setting
 ################################################################
 
 CONST_PORT = 8080
@@ -82,6 +86,8 @@ If you find our work useful for your research or applications, please cite using
 ```
 """
 
+################################################################
+# prepare text examples and image examples
 ################################################################
 
 def get_example_img_list():
@@ -97,6 +103,9 @@ def get_example_txt_list():
 
 example_is = get_example_img_list()
 example_ts = get_example_txt_list()
+
+################################################################
+# initial models
 ################################################################
 
 worker_xbg = Removebg()
@@ -196,8 +205,9 @@ def stage_4_gif(obj_dst, save_folder, do_render_gif=True):
         gif_dst_path = gif_dst
     )
     return gif_dst
-
-#===============================================================
+# ===============================================================
+# gradio display
+# ===============================================================
 with gr.Blocks() as demo:
     gr.Markdown(CONST_HEADER)
     with gr.Row(variant="panel"):
@@ -236,7 +246,12 @@ with gr.Blocks() as demo:
                         imggen_do_render_gif = gr.Checkbox(label="Render gif", value=False, interactive=True)
                         imggen_submit = gr.Button("Generate", variant="primary")       
                     with gr.Row():
-                        gr.Examples(examples=example_is, inputs=[input_image], label="Img examples", examples_per_page=10)
+                        gr.Examples(
+                            examples=example_is, 
+                            inputs=[input_image], 
+                            label="Img examples",
+                            examples_per_page=10
+                        )
            
         with gr.Column(scale=3):
             with gr.Row():
@@ -269,9 +284,12 @@ with gr.Blocks() as demo:
             with gr.Row():    
                 gr.Markdown("""
                 We recommend downloading and opening Glb with 3D software, such as Blender, MeshLab, etc.
+                
                 Limited by gradio, Obj file here only be shown as vertex shading, but Glb can be texture shading.
                 """)
 
+#===============================================================
+# gradio running code
 #===============================================================
 
     none = gr.State(None)
@@ -287,7 +305,9 @@ with gr.Blocks() as demo:
         fn=stage_2_i2v, inputs=[rem_bg_image, textgen_SEED, textgen_STEP, save_folder], 
         outputs=[views_image, cond_image, result_image],
     ).success(
-        fn=stage_3_v23, inputs=[views_image, cond_image, textgen_SEED, save_folder, textgen_max_faces, textgen_do_texture_mapping, textgen_do_render_gif], 
+        fn=stage_3_v23, inputs=[views_image, cond_image, textgen_SEED, save_folder, 
+                                textgen_max_faces, textgen_do_texture_mapping,
+                                textgen_do_render_gif], 
         outputs=[result_3dobj, result_3dglb],
     ).success(
         fn=stage_4_gif, inputs=[result_3dglb, save_folder, textgen_do_render_gif], 
@@ -304,13 +324,17 @@ with gr.Blocks() as demo:
         fn=stage_2_i2v, inputs=[rem_bg_image, imggen_SEED, imggen_STEP, save_folder], 
         outputs=[views_image, cond_image, result_image],
     ).success(
-        fn=stage_3_v23, inputs=[views_image, cond_image, imggen_SEED, save_folder, imggen_max_faces, imggen_do_texture_mapping, imggen_do_render_gif], 
+        fn=stage_3_v23, inputs=[views_image, cond_image, imggen_SEED, save_folder, 
+                                imggen_max_faces, imggen_do_texture_mapping, 
+                                imggen_do_render_gif], 
         outputs=[result_3dobj, result_3dglb],
     ).success(
         fn=stage_4_gif, inputs=[result_3dglb, save_folder, imggen_do_render_gif], 
         outputs=[result_gif],
     ).success(lambda: print('Image_to_3D Done ...'))
     
+#===============================================================
+# start gradio server
 #===============================================================
 
     gr.Markdown(CONST_CITATION)

infer/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,32 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+from .removebg import Removebg
+from .text_to_image import Text2Image
+from .image_to_views import Image2Views, save_gif
+from .views_to_mesh import Views2Mesh
+from .gif_render import GifRenderer
+
+from .utils import seed_everything, auto_amp_inference
+from .utils import get_parameter_number, set_parameter_grad_false

infer/.ipynb_checkpoints/gif_render-checkpoint.py

@@ -0,0 +1,79 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os, sys
+sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}")
+
+from svrm.ldm.vis_util import render
+from infer.utils import seed_everything, timing_decorator
+
+class GifRenderer():
+    '''
+        render frame(s) of mesh using pytorch3d
+    '''
+    def __init__(self, device="cuda:0"):
+        self.device = device
+
+    @timing_decorator("gif render")
+    def __call__(
+        self, 
+        obj_filename, 
+        elev=0, 
+        azim=0, 
+        resolution=512, 
+        gif_dst_path='', 
+        n_views=120, 
+        fps=30, 
+        rgb=True
+    ):
+        render(
+            obj_filename,
+            elev=elev, 
+            azim=azim, 
+            resolution=resolution, 
+            gif_dst_path=gif_dst_path, 
+            n_views=n_views, 
+            fps=fps, 
+            device=self.device, 
+            rgb=rgb
+        )
+
+if __name__ == "__main__":
+    import argparse
+    
+    def get_args():
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--mesh_path", type=str, required=True)
+        parser.add_argument("--output_gif_path", type=str, required=True)
+        parser.add_argument("--device", default="cuda:0", type=str)
+        return parser.parse_args()
+        
+    args = get_args()
+
+    gif_renderer = GifRenderer(device=args.device)
+
+    gif_renderer(
+        args.mesh_path,
+        gif_dst_path = args.output_gif_path
+    )

infer/.ipynb_checkpoints/image_to_views-checkpoint.py

@@ -0,0 +1,126 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os, sys
+sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}")
+
+import time
+import torch
+import random
+import numpy as np
+from PIL import Image
+from einops import rearrange
+from PIL import Image, ImageSequence
+
+from infer.utils import seed_everything, timing_decorator, auto_amp_inference
+from infer.utils import get_parameter_number, set_parameter_grad_false, str_to_bool
+from mvd.hunyuan3d_mvd_std_pipeline import HunYuan3D_MVD_Std_Pipeline
+from mvd.hunyuan3d_mvd_lite_pipeline import Hunyuan3d_MVD_Lite_Pipeline
+
+
+def save_gif(pils, save_path, df=False):
+    # save a list of PIL.Image to gif
+    spf = 4000 / len(pils)
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    pils[0].save(save_path, format="GIF", save_all=True, append_images=pils[1:], duration=spf, loop=0)
+    return save_path
+    
+
+class Image2Views():
+    def __init__(self, device="cuda:0", use_lite=False, save_memory=False):
+        self.device = device
+        if use_lite:
+            self.pipe = Hunyuan3d_MVD_Lite_Pipeline.from_pretrained(
+                "./weights/mvd_lite",
+                torch_dtype = torch.float16,
+                use_safetensors = True,
+            )
+        else:
+            self.pipe = HunYuan3D_MVD_Std_Pipeline.from_pretrained(
+                "./weights/mvd_std",
+                torch_dtype = torch.float16,
+                use_safetensors = True,
+            )
+        self.pipe = self.pipe.to(device)
+        self.order = [0, 1, 2, 3, 4, 5] if use_lite else [0, 2, 4, 5, 3, 1]
+        self.save_memory = save_memory
+        set_parameter_grad_false(self.pipe.unet)
+        print('image2views unet model', get_parameter_number(self.pipe.unet))
+
+    @torch.no_grad()
+    @timing_decorator("image to views")
+    @auto_amp_inference
+    def __call__(self, *args, **kwargs):
+        if self.save_memory:
+            self.pipe = self.pipe.to(self.device)
+            torch.cuda.empty_cache()
+            res = self.call(*args, **kwargs)
+            self.pipe = self.pipe.to("cpu")
+        else:
+            res = self.call(*args, **kwargs)
+        torch.cuda.empty_cache()
+        return res
+        
+    def call(self, pil_img, seed=0, steps=50, guidance_scale=2.0):
+        seed_everything(seed)
+        generator = torch.Generator(device=self.device)
+        res_img = self.pipe(pil_img, 
+                            num_inference_steps=steps,
+                            guidance_scale=guidance_scale, 
+                            generat=generator).images
+        show_image = rearrange(np.asarray(res_img[0], dtype=np.uint8), '(n h) (m w) c -> (n m) h w c', n=3, m=2)
+        pils = [res_img[1]]+[Image.fromarray(show_image[idx]) for idx in self.order] 
+        torch.cuda.empty_cache()
+        return res_img, pils
+
+
+if __name__ == "__main__":
+    import argparse
+    
+    def get_args():
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--rgba_path", type=str, required=True)
+        parser.add_argument("--output_views_path", type=str, required=True)
+        parser.add_argument("--output_cond_path", type=str, required=True)
+        parser.add_argument("--seed", default=0, type=int)
+        parser.add_argument("--steps", default=50, type=int)
+        parser.add_argument("--device", default="cuda:0", type=str)
+        parser.add_argument("--use_lite", default='false', type=str)
+        return parser.parse_args()
+        
+    args = get_args()
+
+    args.use_lite = str_to_bool(args.use_lite)
+
+    rgba_pil = Image.open(args.rgba_path)
+
+    assert rgba_pil.mode == "RGBA", "rgba_pil must be RGBA mode"
+
+    model = Image2Views(device=args.device, use_lite=args.use_lite)
+
+    (views_pil, cond), _ = model(rgba_pil, seed=args.seed, steps=args.steps)
+
+    views_pil.save(args.output_views_path)
+    cond.save(args.output_cond_path)
+

infer/.ipynb_checkpoints/removebg-checkpoint.py

@@ -0,0 +1,101 @@
+import os, sys
+sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}")
+
+import numpy as np
+from PIL import Image
+from rembg import remove, new_session
+from infer.utils import timing_decorator
+
+class Removebg():
+    def __init__(self, name="u2net"):
+        self.session = new_session(name)
+
+    @timing_decorator("remove background")
+    def __call__(self, rgb_maybe, force=True):
+        '''
+            args:
+                rgb_maybe: PIL.Image, with RGB mode or RGBA mode
+                force: bool, if input is RGBA mode, covert to RGB then remove bg
+            return:
+                rgba_img: PIL.Image, with RGBA mode
+        '''
+        if rgb_maybe.mode == "RGBA":
+            if force:  
+                rgb_maybe = rgb_maybe.convert("RGB")
+                rgba_img = remove(rgb_maybe, session=self.session)
+            else:
+                rgba_img = rgb_maybe
+        else:
+            rgba_img = remove(rgb_maybe, session=self.session)
+            
+        rgba_img = white_out_background(rgba_img)
+        
+        rgba_img = preprocess(rgba_img)
+        
+        return rgba_img
+
+
+def white_out_background(pil_img):
+    data = pil_img.getdata()
+    new_data = []
+    for r, g, b, a in data:
+        if a < 16:  # background
+            new_data.append((255, 255, 255, 0))  # full white color
+        else:
+            is_white = (r>235) and (g>235) and (b>235)
+            new_r = 235 if is_white else r
+            new_g = 235 if is_white else g
+            new_b = 235 if is_white else b
+            new_data.append((new_r, new_g, new_b, a))
+    pil_img.putdata(new_data)
+    return pil_img
+    
+def preprocess(rgba_img, size=(512,512), ratio=1.15):
+    image = np.asarray(rgba_img)
+    rgb, alpha = image[:,:,:3] / 255., image[:,:,3:] / 255.
+
+    # crop
+    coords = np.nonzero(alpha > 0.1)
+    x_min, x_max = coords[0].min(), coords[0].max()
+    y_min, y_max = coords[1].min(), coords[1].max()
+    rgb = (rgb[x_min:x_max, y_min:y_max, :] * 255).astype("uint8")
+    alpha = (alpha[x_min:x_max, y_min:y_max, 0] * 255).astype("uint8")
+
+    # padding
+    h, w = rgb.shape[:2]
+    resize_side = int(max(h, w) * ratio)
+    pad_h, pad_w = resize_side - h, resize_side - w
+    start_h, start_w = pad_h // 2, pad_w // 2
+    new_rgb = np.ones((resize_side, resize_side, 3), dtype=np.uint8) * 255
+    new_alpha = np.zeros((resize_side, resize_side), dtype=np.uint8)
+    new_rgb[start_h:start_h + h, start_w:start_w + w] = rgb
+    new_alpha[start_h:start_h + h, start_w:start_w + w] = alpha
+    rgba_array = np.concatenate((new_rgb, new_alpha[:,:,None]), axis=-1)
+    
+    rgba_image = Image.fromarray(rgba_array, 'RGBA')
+    rgba_image = rgba_image.resize(size)
+    return rgba_image
+
+
+if __name__ == "__main__":
+    
+    import argparse
+    
+    def get_args():
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--rgb_path", type=str, required=True)
+        parser.add_argument("--output_rgba_path", type=str, required=True)
+        parser.add_argument("--force", default=False, action="store_true")
+        return parser.parse_args()
+        
+    args = get_args()
+
+    rgb_maybe = Image.open(args.rgb_path)
+    
+    model = Removebg()
+
+    rgba_pil = model(rgb_maybe, args.force)
+
+    rgba_pil.save(args.output_rgba_path)
+
+    

infer/.ipynb_checkpoints/text_to_image-checkpoint.py

@@ -0,0 +1,105 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+import os , sys
+sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}")
+
+import torch
+from diffusers import HunyuanDiTPipeline, AutoPipelineForText2Image
+
+from infer.utils import seed_everything, timing_decorator, auto_amp_inference
+from infer.utils import get_parameter_number, set_parameter_grad_false
+
+
+class Text2Image():
+    def __init__(self, pretrain="weights/hunyuanDiT", device="cuda:0", save_memory=None):
+        '''
+            save_memory: if GPU memory is low, can set it
+        '''
+        self.save_memory = save_memory
+        self.device = device
+        self.pipe = AutoPipelineForText2Image.from_pretrained(
+            pretrain, 
+            torch_dtype = torch.float16, 
+            enable_pag = True, 
+            pag_applied_layers = ["blocks.(16|17|18|19)"]
+        )
+        set_parameter_grad_false(self.pipe.transformer)
+        print('text2image transformer model', get_parameter_number(self.pipe.transformer))
+        if not save_memory: 
+            self.pipe = self.pipe.to(device)
+        self.neg_txt = "文本,特写,裁剪,出框,最差质量,低质量,JPEG伪影,PGLY,重复,病态,残缺,多余的手指,变异的手," \
+                       "画得不好的手,画得不好的脸,变异,畸形,模糊,脱水,糟糕的解剖学,糟糕的比例,多余的肢体,克隆的脸," \
+                       "毁容,恶心的比例,畸形的肢体,缺失的手臂,缺失的腿,额外的手臂,额外的腿,融合的手指,手指太多,长脖子"
+
+    @torch.no_grad()
+    @timing_decorator('text to image')
+    @auto_amp_inference
+    def __call__(self, *args, **kwargs):
+        if self.save_memory:
+            self.pipe = self.pipe.to(self.device)
+            torch.cuda.empty_cache()
+            res = self.call(*args, **kwargs)
+            self.pipe = self.pipe.to("cpu")
+        else:
+            res = self.call(*args, **kwargs)
+        torch.cuda.empty_cache()
+        return res
+
+    def call(self, prompt, seed=0, steps=25):
+        '''
+            args:
+                prompr: str
+                seed: int
+                steps: int
+            return:
+                rgb: PIL.Image
+        '''
+        print("prompt is:", prompt)
+        prompt = prompt + ",白色背景,3D风格,最佳质量"
+        seed_everything(seed)
+        generator = torch.Generator(device=self.device)
+        if seed is not None: generator = generator.manual_seed(int(seed))
+        rgb = self.pipe(prompt=prompt, negative_prompt=self.neg_txt, num_inference_steps=steps, 
+            pag_scale=1.3, width=1024, height=1024, generator=generator, return_dict=False)[0][0]
+        torch.cuda.empty_cache()
+        return rgb
+
+if __name__ == "__main__":
+    import argparse
+    
+    def get_args():
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--text2image_path", default="weights/hunyuanDiT", type=str)
+        parser.add_argument("--text_prompt", default="", type=str)
+        parser.add_argument("--output_img_path", default="./outputs/test/img.jpg", type=str)
+        parser.add_argument("--device", default="cuda:0", type=str)
+        parser.add_argument("--seed", default=0, type=int)
+        parser.add_argument("--steps", default=25, type=int)
+        return parser.parse_args()
+    args = get_args()
+    
+    text2image_model = Text2Image(device=args.device)
+    rgb_img = text2image_model(args.text_prompt, seed=args.seed, steps=args.steps)
+    rgb_img.save(args.output_img_path)
+    

infer/.ipynb_checkpoints/utils-checkpoint.py

@@ -0,0 +1,87 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os
+import time
+import random
+import numpy as np
+import torch
+from torch.cuda.amp import autocast, GradScaler
+from functools import wraps
+
+def seed_everything(seed):
+    '''
+        seed everthing
+    '''
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    os.environ["PL_GLOBAL_SEED"] = str(seed)
+
+def timing_decorator(category: str):
+    '''
+        timing_decorator: record time
+    '''
+    def decorator(func):
+        func.call_count = 0
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs) 
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            func.call_count += 1
+            print(f"[HunYuan3D]-[{category}], cost time: {elapsed_time:.4f}s") # huiwen
+            return result
+        return wrapper
+    return decorator
+
+def auto_amp_inference(func):
+    '''
+        with torch.cuda.amp.autocast()"
+            xxx
+    '''
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with autocast(): 
+            output = func(*args, **kwargs) 
+        return output
+    return wrapper
+
+def get_parameter_number(model):
+    total_num = sum(p.numel() for p in model.parameters())
+    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return {'Total': total_num, 'Trainable': trainable_num}
+
+def set_parameter_grad_false(model):
+    for p in model.parameters():
+        p.requires_grad = False
+
+def str_to_bool(s):
+    if s.lower() in ['true', 't', 'yes', 'y', '1']: 
+        return True
+    elif s.lower() in ['false', 'f', 'no', 'n', '0']: 
+        return False
+    else: 
+        raise f"bool arg must one of ['true', 't', 'yes', 'y', '1', 'false', 'f', 'no', 'n', '0']"

infer/.ipynb_checkpoints/views_to_mesh-checkpoint.py

@@ -0,0 +1,154 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os, sys
+sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}")
+
+import time
+import torch
+import random
+import numpy as np
+from PIL import Image
+from einops import rearrange
+from PIL import Image, ImageSequence
+
+from infer.utils import seed_everything, timing_decorator, auto_amp_inference
+from infer.utils import get_parameter_number, set_parameter_grad_false, str_to_bool
+from svrm.predictor import MV23DPredictor
+
+
+class Views2Mesh():
+    def __init__(self, mv23d_cfg_path, mv23d_ckt_path, 
+                 device="cuda:0", use_lite=False, save_memory=False):
+        '''
+            mv23d_cfg_path: config yaml file 
+            mv23d_ckt_path: path to ckpt
+            use_lite: lite version
+            save_memory: cpu auto
+        '''
+        self.mv23d_predictor = MV23DPredictor(mv23d_ckt_path, mv23d_cfg_path, device=device)  
+        self.mv23d_predictor.model.eval()
+        self.order = [0, 1, 2, 3, 4, 5] if use_lite else [0, 2, 4, 5, 3, 1]
+        self.device = device
+        self.save_memory = save_memory
+        set_parameter_grad_false(self.mv23d_predictor.model)
+        print('view2mesh model', get_parameter_number(self.mv23d_predictor.model))
+
+    @torch.no_grad()
+    @timing_decorator("views to mesh")
+    @auto_amp_inference
+    def __call__(self, *args, **kwargs):
+        if self.save_memory:
+            self.mv23d_predictor.model = self.mv23d_predictor.model.to(self.device)
+            torch.cuda.empty_cache()
+            res = self.call(*args, **kwargs)
+            self.mv23d_predictor.model = self.mv23d_predictor.model.to("cpu")
+        else:
+            res = self.call(*args, **kwargs)
+        torch.cuda.empty_cache()
+        return res
+
+    def call(
+        self,
+        views_pil=None, 
+        cond_pil=None, 
+        gif_pil=None, 
+        seed=0, 
+        target_face_count = 10000,
+        do_texture_mapping = True,
+        save_folder='./outputs/test'
+    ):
+        '''
+            can set views_pil, cond_pil simutaously or set gif_pil only
+            seed: int
+            target_face_count: int 
+            save_folder: path to save mesh files
+        '''
+        save_dir = save_folder
+        os.makedirs(save_dir, exist_ok=True)
+
+        if views_pil is not None and cond_pil is not None:
+            show_image = rearrange(np.asarray(views_pil, dtype=np.uint8), 
+                                   '(n h) (m w) c -> (n m) h w c', n=3, m=2)
+            views = [Image.fromarray(show_image[idx]) for idx in self.order] 
+            image_list = [cond_pil]+ views
+            image_list = [img.convert('RGB') for img in image_list]
+        elif gif_pil is not None:
+            image_list = [img.convert('RGB') for img in ImageSequence.Iterator(gif_pil)]
+        
+        image_input = image_list[0]
+        image_list = image_list[1:] + image_list[:1]
+        
+        seed_everything(seed)
+        self.mv23d_predictor.predict(
+            image_list, 
+            save_dir = save_dir, 
+            image_input = image_input,
+            target_face_count = target_face_count,
+            do_texture_mapping = do_texture_mapping
+        )
+        torch.cuda.empty_cache()
+        return save_dir
+
+
+if __name__ == "__main__":
+    
+    import argparse
+    
+    def get_args():
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--views_path", type=str, required=True)
+        parser.add_argument("--cond_path", type=str, required=True)
+        parser.add_argument("--save_folder", default="./outputs/test/", type=str)
+        parser.add_argument("--mv23d_cfg_path", default="./svrm/configs/svrm.yaml", type=str)
+        parser.add_argument("--mv23d_ckt_path", default="weights/svrm/svrm.safetensors", type=str)
+        parser.add_argument("--max_faces_num", default=90000, type=int, 
+            help="max num of face, suggest 90000 for effect, 10000 for speed")
+        parser.add_argument("--device", default="cuda:0", type=str)
+        parser.add_argument("--use_lite", default='false', type=str)
+        parser.add_argument("--do_texture_mapping", default='false', type=str)
+        
+        return parser.parse_args()
+        
+    args = get_args()
+    args.use_lite = str_to_bool(args.use_lite)
+    args.do_texture_mapping = str_to_bool(args.do_texture_mapping)
+
+    views = Image.open(args.views_path)
+    cond = Image.open(args.cond_path)
+
+    views_to_mesh_model = Views2Mesh(
+        args.mv23d_cfg_path, 
+        args.mv23d_ckt_path,
+        device = args.device, 
+        use_lite = args.use_lite
+    )
+
+    views_to_mesh_model(
+        views,  cond,  0,
+        target_face_count = args.max_faces_num,
+        save_folder = args.save_folder,
+        do_texture_mapping = args.do_texture_mapping
+    )
+

infer/__init__.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

infer/gif_render.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

infer/image_to_views.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

infer/text_to_image.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

infer/utils.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

infer/views_to_mesh.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

main.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

mvd/.ipynb_checkpoints/hunyuan3d_mvd_lite_pipeline-checkpoint.py

@@ -0,0 +1,392 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import math
+import numpy
+import torch
+import inspect
+import warnings
+from PIL import Image
+from einops import rearrange
+import torch.nn.functional as F
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from typing import Any, Callable, Dict, List, Optional, Union
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers import DDPMScheduler, EulerAncestralDiscreteScheduler, ImagePipelineOutput
+from diffusers.loaders import (
+    FromSingleFileMixin, 
+    LoraLoaderMixin, 
+    TextualInversionLoaderMixin
+)
+from transformers import (
+    CLIPImageProcessor, 
+    CLIPTextModel, 
+    CLIPTokenizer, 
+    CLIPVisionModelWithProjection
+)
+from diffusers.models.attention_processor import (
+    Attention, 
+    AttnProcessor, 
+    XFormersAttnProcessor, 
+    AttnProcessor2_0
+)
+
+from .utils import to_rgb_image, white_out_background, recenter_img
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from here import Hunyuan3d_MVD_Lite_Pipeline
+
+        >>> pipe = Hunyuan3d_MVD_Lite_Pipeline.from_pretrained(
+        ...     "weights/mvd_lite", torch_dtype=torch.float16
+        ... )
+        >>> pipe.to("cuda")
+
+        >>> img = Image.open("demo.png")
+        >>> res_img = pipe(img).images[0]
+"""
+
+def unscale_latents(latents): return latents / 0.75 + 0.22
+def unscale_image  (image  ): return   image / 0.50 * 0.80
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+
+class ReferenceOnlyAttnProc(torch.nn.Module):
+    # reference attention
+    def __init__(self, chained_proc, enabled=False, name=None):
+        super().__init__()
+        self.enabled = enabled
+        self.chained_proc = chained_proc
+        self.name = name
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, mode="w", ref_dict=None):
+        if encoder_hidden_states is None: encoder_hidden_states = hidden_states
+        if self.enabled:
+            if mode == 'w': 
+                ref_dict[self.name] = encoder_hidden_states
+            elif mode == 'r': 
+                encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1)
+        res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask)
+        return res
+
+
+class RefOnlyNoisedUNet(torch.nn.Module):
+    def __init__(self, unet, train_sched, val_sched):
+        super().__init__()
+        self.unet = unet
+        self.train_sched = train_sched
+        self.val_sched = val_sched
+
+        unet_lora_attn_procs = dict()
+        for name, _ in unet.attn_processors.items():
+            unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(AttnProcessor2_0(), 
+                                                           enabled=name.endswith("attn1.processor"), 
+                                                           name=name)
+        unet.set_attn_processor(unet_lora_attn_procs)
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.unet, name)
+
+    def forward(self, sample, timestep, encoder_hidden_states, *args, cross_attention_kwargs, **kwargs):
+        cond_lat = cross_attention_kwargs['cond_lat']
+        noise = torch.randn_like(cond_lat)
+        if self.training:
+            noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)
+            noisy_cond_lat = self.train_sched.scale_model_input(noisy_cond_lat, timestep)
+        else:
+            noisy_cond_lat = self.val_sched.add_noise(cond_lat, noise, timestep.reshape(-1))
+            noisy_cond_lat = self.val_sched.scale_model_input(noisy_cond_lat, timestep.reshape(-1))
+
+        ref_dict = {}
+        self.unet(noisy_cond_lat, 
+                  timestep, 
+                  encoder_hidden_states, 
+                  *args, 
+                  cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict), 
+                  **kwargs)
+        return  self.unet(sample, 
+                          timestep, 
+                          encoder_hidden_states, 
+                          *args, 
+                          cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict), 
+                          **kwargs)
+
+
+class Hunyuan3d_MVD_Lite_Pipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        vision_encoder: CLIPVisionModelWithProjection,
+        feature_extractor_clip: CLIPImageProcessor, 
+        feature_extractor_vae: CLIPImageProcessor,
+        ramping_coefficients: Optional[list] = None,
+        safety_checker=None,
+    ):
+        DiffusionPipeline.__init__(self)
+        self.register_modules(
+            vae=vae,
+            unet=unet,
+            tokenizer=tokenizer,
+            scheduler=scheduler,
+            text_encoder=text_encoder,
+            vision_encoder=vision_encoder,
+            feature_extractor_vae=feature_extractor_vae,
+            feature_extractor_clip=feature_extractor_clip
+        )
+            # rewrite the stable diffusion pipeline
+            # vae: vae
+            # unet: unet
+            # tokenizer: tokenizer
+            # scheduler: scheduler
+            # text_encoder: text_encoder
+            # vision_encoder: vision_encoder
+            # feature_extractor_vae: feature_extractor_vae
+            # feature_extractor_clip: feature_extractor_clip
+        self.register_to_config(ramping_coefficients=ramping_coefficients)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        extra_step_kwargs = {}
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_eta: extra_step_kwargs["eta"] = eta
+
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator: extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @torch.no_grad()
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None: uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError()
+            elif isinstance(negative_prompt, str): uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt): raise ValueError()
+            else: uncond_tokens = negative_prompt
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(uncond_tokens, 
+                                          padding="max_length", 
+                                          max_length=max_length, 
+                                          truncation=True, 
+                                          return_tensors="pt")
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(device), attention_mask=attention_mask)
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            seq_len = negative_prompt_embeds.shape[1]
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+    @torch.no_grad()
+    def encode_condition_image(self, image: torch.Tensor): return self.vae.encode(image).latent_dist.sample()
+
+    @torch.no_grad()
+    def __call__(self, image=None, 
+                 width=640, 
+                 height=960, 
+                 num_inference_steps=75, 
+                 return_dict=True, 
+                 generator=None, 
+                 **kwargs):
+        batch_size = 1
+        num_images_per_prompt = 1
+        output_type = 'pil'
+        do_classifier_free_guidance = True
+        guidance_rescale = 0.
+        if isinstance(self.unet, UNet2DConditionModel): 
+            self.unet = RefOnlyNoisedUNet(self.unet, None, self.scheduler).eval()
+
+        cond_image = recenter_img(image)
+        cond_image = to_rgb_image(image)
+        image = cond_image
+        image_1 = self.feature_extractor_vae(images=image, return_tensors="pt").pixel_values
+        image_2 = self.feature_extractor_clip(images=image, return_tensors="pt").pixel_values
+        image_1 = image_1.to(device=self.vae.device, dtype=self.vae.dtype)
+        image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)
+
+        cond_lat = self.encode_condition_image(image_1)
+        negative_lat = self.encode_condition_image(torch.zeros_like(image_1))
+        cond_lat = torch.cat([negative_lat, cond_lat])
+        cross_attention_kwargs = dict(cond_lat=cond_lat)
+
+        global_embeds = self.vision_encoder(image_2, output_hidden_states=False).image_embeds.unsqueeze(-2)
+        encoder_hidden_states = self._encode_prompt('', self.device, num_images_per_prompt, False)
+        ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        prompt_embeds = torch.cat([encoder_hidden_states, encoder_hidden_states + global_embeds * ramp])
+
+        device = self._execution_device
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(batch_size * num_images_per_prompt,
+                                       num_channels_latents, 
+                                       height, 
+                                       width, 
+                                       prompt_embeds.dtype, 
+                                       device, 
+                                       generator, 
+                                       None)
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, 0.0)
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        # set adaptive cfg
+        # the image order is:
+        #    [0, 60, 
+        #     120, 180,
+        #     240, 300]
+        # the cfg is set as 3, 2.5, 2, 1.5
+        
+        tmp_guidance_scale = torch.ones_like(latents)
+        tmp_guidance_scale[:, :, :40, :40] = 3
+        tmp_guidance_scale[:, :, :40, 40:] =  2.5
+        tmp_guidance_scale[:, :, 40:80, :40] =  2
+        tmp_guidance_scale[:, :, 40:80, 40:] =  1.5
+        tmp_guidance_scale[:, :, 80:120, :40] =  2
+        tmp_guidance_scale[:, :, 80:120, 40:] =  2.5
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                noise_pred = self.unet(latent_model_input, t,
+                                encoder_hidden_states=prompt_embeds, 
+                                cross_attention_kwargs=cross_attention_kwargs, 
+                                return_dict=False)[0]
+
+                adaptive_guidance_scale = (2 + 16 * (t / 1000) ** 5) / 3
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + \
+                        tmp_guidance_scale * adaptive_guidance_scale * \
+                        (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                if i==len(timesteps)-1 or ((i+1)>num_warmup_steps and (i+1)%self.scheduler.order==0): 
+                    progress_bar.update()
+
+        latents = unscale_latents(latents)
+        image = unscale_image(self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0])
+        image = self.image_processor.postprocess(image, output_type='pil')[0]
+        image = [image, cond_image]
+        return ImagePipelineOutput(images=image) if return_dict else (image,)
+

mvd/.ipynb_checkpoints/hunyuan3d_mvd_std_pipeline-checkpoint.py

@@ -0,0 +1,473 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import inspect
+from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import os
+import torch
+import numpy as np
+from PIL import Image
+
+import diffusers
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.models.attention_processor import (
+    Attention, 
+    AttnProcessor, 
+    XFormersAttnProcessor, 
+    AttnProcessor2_0
+)
+from diffusers import (
+    AutoencoderKL, 
+    DDPMScheduler,
+    DiffusionPipeline, 
+    EulerAncestralDiscreteScheduler, 
+    UNet2DConditionModel, 
+    ImagePipelineOutput
+)
+import transformers
+from transformers import (
+    CLIPImageProcessor, 
+    CLIPTextModel, 
+    CLIPTokenizer, 
+    CLIPVisionModelWithProjection, 
+    CLIPTextModelWithProjection
+)
+
+from .utils import to_rgb_image, white_out_background, recenter_img
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import Hunyuan3d_MVD_XL_Pipeline
+
+        >>> pipe = Hunyuan3d_MVD_XL_Pipeline.from_pretrained(
+        ...     "Tencent-Hunyuan-3D/MVD-XL", torch_dtype=torch.float16
+        ... )
+        >>> pipe.to("cuda")
+
+        >>> img = Image.open("demo.png")
+        >>> res_img = pipe(img).images[0]
+        ```
+"""
+
+
+
+def scale_latents(latents):   return (latents - 0.22) * 0.75
+def unscale_latents(latents): return (latents / 0.75) + 0.22
+def scale_image(image):       return (image - 0.5) / 0.5
+def scale_image_2(image):     return (image * 0.5) / 0.8
+def unscale_image(image):     return (image * 0.5) + 0.5
+def unscale_image_2(image):   return (image * 0.8) / 0.5
+
+
+
+
+class ReferenceOnlyAttnProc(torch.nn.Module):
+    def __init__(self, chained_proc, enabled=False, name=None):
+        super().__init__()
+        self.enabled = enabled
+        self.chained_proc = chained_proc
+        self.name = name
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, mode="w", ref_dict=None):
+        encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states
+        if self.enabled:
+            if   mode == 'w': ref_dict[self.name]   = encoder_hidden_states
+            elif mode == 'r': encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1)
+            else:             raise Exception(f"mode should not be {mode}")
+        return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask)
+
+
+class RefOnlyNoisedUNet(torch.nn.Module):
+    def __init__(self, unet, scheduler) -> None:
+        super().__init__()
+        self.unet = unet
+        self.scheduler = scheduler
+
+        unet_attn_procs = dict()
+        for name, _ in unet.attn_processors.items():
+            if torch.__version__ >= '2.0': default_attn_proc = AttnProcessor2_0()
+            elif is_xformers_available():  default_attn_proc = XFormersAttnProcessor()
+            else:                          default_attn_proc = AttnProcessor()
+            unet_attn_procs[name] = ReferenceOnlyAttnProc(
+                default_attn_proc, enabled=name.endswith("attn1.processor"), name=name
+            )
+        unet.set_attn_processor(unet_attn_procs)
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.unet, name)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        class_labels: Optional[torch.Tensor] = None,
+        down_block_res_samples: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_res_sample: Optional[Tuple[torch.Tensor]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        return_dict: bool = True,
+        **kwargs
+    ):
+
+        dtype = self.unet.dtype
+
+        # cond_lat add same level noise
+        cond_lat = cross_attention_kwargs['cond_lat']
+        noise = torch.randn_like(cond_lat)
+
+        noisy_cond_lat = self.scheduler.add_noise(cond_lat, noise, timestep.reshape(-1))
+        noisy_cond_lat = self.scheduler.scale_model_input(noisy_cond_lat, timestep.reshape(-1))
+
+        ref_dict = {}
+        
+        _ = self.unet(
+            noisy_cond_lat, 
+            timestep, 
+            encoder_hidden_states = encoder_hidden_states, 
+            class_labels = class_labels,
+            cross_attention_kwargs = dict(mode="w", ref_dict=ref_dict),
+            added_cond_kwargs = added_cond_kwargs,
+            return_dict = return_dict,
+            **kwargs
+        )
+
+        res = self.unet(
+            sample, 
+            timestep, 
+            encoder_hidden_states, 
+            class_labels=class_labels,
+            cross_attention_kwargs = dict(mode="r", ref_dict=ref_dict),
+            down_block_additional_residuals = [
+                sample.to(dtype=dtype) for sample in down_block_res_samples
+            ] if down_block_res_samples is not None else None,
+            mid_block_additional_residual = (
+                mid_block_res_sample.to(dtype=dtype) 
+                if mid_block_res_sample is not None else None),
+            added_cond_kwargs = added_cond_kwargs,
+            return_dict = return_dict,
+            **kwargs
+        )
+        return res
+        
+
+
+class HunYuan3D_MVD_Std_Pipeline(diffusers.DiffusionPipeline):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        feature_extractor_vae: CLIPImageProcessor,
+        vision_processor: CLIPImageProcessor,
+        vision_encoder: CLIPVisionModelWithProjection,
+        vision_encoder_2: CLIPVisionModelWithProjection,
+        ramping_coefficients: Optional[list] = None,
+        add_watermarker: Optional[bool] = None,
+        safety_checker = None,
+    ):
+        DiffusionPipeline.__init__(self)
+
+        self.register_modules(
+            vae=vae, unet=unet, scheduler=scheduler, safety_checker=None, feature_extractor_vae=feature_extractor_vae,
+            vision_processor=vision_processor, vision_encoder=vision_encoder, vision_encoder_2=vision_encoder_2, 
+        )
+        self.register_to_config( ramping_coefficients = ramping_coefficients)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.default_sample_size = self.unet.config.sample_size 
+        self.watermark = None 
+        self.prepare_init = False
+
+    def prepare(self):
+        assert isinstance(self.unet, UNet2DConditionModel), "unet should be UNet2DConditionModel"
+        self.unet = RefOnlyNoisedUNet(self.unet, self.scheduler).eval()
+        self.prepare_init = True
+
+    def encode_image(self, image: torch.Tensor, scale_factor: bool = False):
+        latent = self.vae.encode(image).latent_dist.sample()
+        return (latent * self.vae.config.scaling_factor) if scale_factor else latent
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) 
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def _get_add_time_ids(
+        self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None
+    ):
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+
+        passed_add_embed_dim = (
+            self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+        )
+        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+        if expected_add_embed_dim != passed_add_embed_dim:
+            raise ValueError(
+                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, " \
+                f"but a vector of {passed_add_embed_dim} was created. The model has an incorrect config." \
+                f" Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+            )
+
+        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+        return add_time_ids
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:  extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator: extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
+        
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Image.Image = None,
+        guidance_scale = 2.0,
+        output_type: Optional[str] = "pil",
+        num_inference_steps: int = 50,
+        return_dict: bool = True,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        latent: torch.Tensor = None,
+        guidance_curve = None,
+        **kwargs
+    ):
+        if not self.prepare_init:
+            self.prepare()
+
+        here = dict(device=self.vae.device, dtype=self.vae.dtype)
+            
+        batch_size = 1
+        num_images_per_prompt = 1
+        width, height = 512 * 2,  512 * 3
+        target_size = original_size = (height, width)
+
+        self._guidance_scale = guidance_scale
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        device = self._execution_device
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        
+        # Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            self.vae.dtype,
+            device,
+            generator,
+            latents=latent,
+        )
+        
+        # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+
+        # Prepare added time ids & embeddings
+        text_encoder_projection_dim = 1280     
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype=self.vae.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        negative_add_time_ids = add_time_ids
+
+        # hw: preprocess
+        cond_image = recenter_img(image)
+        cond_image = to_rgb_image(image)
+        image_vae = self.feature_extractor_vae(images=cond_image, return_tensors="pt").pixel_values.to(**here)
+        image_clip = self.vision_processor(images=cond_image, return_tensors="pt").pixel_values.to(**here)
+
+        # hw: get cond_lat from cond_img using vae
+        cond_lat = self.encode_image(image_vae, scale_factor=False)
+        negative_lat = self.encode_image(torch.zeros_like(image_vae), scale_factor=False) 
+        cond_lat = torch.cat([negative_lat, cond_lat])
+
+        # hw: get visual global embedding using clip
+        global_embeds_1 = self.vision_encoder(image_clip, output_hidden_states=False).image_embeds.unsqueeze(-2)
+        global_embeds_2 = self.vision_encoder_2(image_clip, output_hidden_states=False).image_embeds.unsqueeze(-2)
+        global_embeds = torch.concat([global_embeds_1, global_embeds_2], dim=-1)
+        
+        ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
+        prompt_embeds = self.uc_text_emb.to(**here)
+        pooled_prompt_embeds =  self.uc_text_emb_2.to(**here)
+        
+        prompt_embeds = prompt_embeds + global_embeds * ramp
+        add_text_embeds = pooled_prompt_embeds
+        
+        if self.do_classifier_free_guidance:
+            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        # Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        timestep_cond = None
+        self._num_timesteps = len(timesteps)
+
+        if guidance_curve is None:
+            guidance_curve = lambda t: guidance_scale
+        
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                    
+                 # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # predict the noise residual
+                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                    
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=dict(cond_lat=cond_lat),
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                
+                # cur_guidance_scale = self.guidance_scale
+                cur_guidance_scale = guidance_curve(t)  # 1.5 + 2.5 * ((t/1000)**2)
+                
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + cur_guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                    # cur_guidance_scale_topleft = (cur_guidance_scale - 1.0) * 4 + 1.0
+                    # noise_pred_top_left = noise_pred_uncond + 
+                    #    cur_guidance_scale_topleft * (noise_pred_text - noise_pred_uncond)
+                    # _, _, h, w = noise_pred.shape
+                    # noise_pred[:, :, :h//3, :w//2] = noise_pred_top_left[:, :, :h//3, :w//2]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+        
+        latents = unscale_latents(latents)
+
+        if output_type=="latent":
+            image = latents
+        else:
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            image = unscale_image(unscale_image_2(image)).clamp(0, 1)
+            image = [
+                Image.fromarray((image[0]*255+0.5).clamp_(0, 255).permute(1, 2, 0).cpu().numpy().astype("uint8")),
+                # self.image_processor.postprocess(image, output_type=output_type)[0],
+                cond_image.resize((512, 512))
+            ]
+
+        if not return_dict: return (image,)
+        return ImagePipelineOutput(images=image)
+
+    def save_pretrained(self, save_directory):
+        # uc_text_emb.pt and uc_text_emb_2.pt are inferenced and saved in advance
+        super().save_pretrained(save_directory)
+        torch.save(self.uc_text_emb, os.path.join(save_directory, "uc_text_emb.pt"))
+        torch.save(self.uc_text_emb_2, os.path.join(save_directory, "uc_text_emb_2.pt"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # uc_text_emb.pt and uc_text_emb_2.pt are inferenced and saved in advance
+        pipeline = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        pipeline.uc_text_emb = torch.load(os.path.join(pretrained_model_name_or_path, "uc_text_emb.pt"))
+        pipeline.uc_text_emb_2 = torch.load(os.path.join(pretrained_model_name_or_path, "uc_text_emb_2.pt"))
+        return pipeline

mvd/.ipynb_checkpoints/utils-checkpoint.py

@@ -0,0 +1,87 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import numpy as np
+from PIL import Image
+
+def to_rgb_image(maybe_rgba: Image.Image):
+    '''
+        convert a PIL.Image to rgb mode with white background
+        maybe_rgba: PIL.Image
+        return: PIL.Image
+    '''
+    if maybe_rgba.mode == 'RGB':
+        return maybe_rgba
+    elif maybe_rgba.mode == 'RGBA':
+        rgba = maybe_rgba
+        img = np.random.randint(255, 256, size=[rgba.size[1], rgba.size[0], 3], dtype=np.uint8)
+        img = Image.fromarray(img, 'RGB')
+        img.paste(rgba, mask=rgba.getchannel('A'))
+        return img
+    else:
+        raise ValueError("Unsupported image type.", maybe_rgba.mode)
+        
+def white_out_background(pil_img, is_gray_fg=True):
+    data = pil_img.getdata()
+    new_data = []
+    #  convert fore-ground white to gray
+    for r, g, b, a in data:
+        if a < 16: 
+            new_data.append((255, 255, 255, 0))  # back-ground to be black
+        else:
+            is_white = is_gray_fg and (r>235) and (g>235) and (b>235)
+            new_r = 235 if is_white else r
+            new_g = 235 if is_white else g
+            new_b = 235 if is_white else b
+            new_data.append((new_r, new_g, new_b, a))
+    pil_img.putdata(new_data)
+    return pil_img
+    
+def recenter_img(img, size=512, color=(255,255,255)):
+    img = white_out_background(img)
+    mask = np.array(img)[..., 3]
+    image = np.array(img)[..., :3]
+    
+    H, W, C = image.shape
+    coords = np.nonzero(mask)
+    x_min, x_max = coords[0].min(), coords[0].max()
+    y_min, y_max = coords[1].min(), coords[1].max()
+    h = x_max - x_min
+    w = y_max - y_min
+    if h == 0 or w == 0: raise ValueError
+    roi = image[x_min:x_max, y_min:y_max]
+
+    border_ratio = 0.15 # 0.2
+    pad_h = int(h * border_ratio)
+    pad_w = int(w * border_ratio)
+
+    result_tmp = np.full((h + pad_h, w + pad_w, C), color, dtype=np.uint8)
+    result_tmp[pad_h // 2: pad_h // 2 + h, pad_w // 2: pad_w // 2 + w] = roi
+
+    cur_h, cur_w = result_tmp.shape[:2]
+    side = max(cur_h, cur_w)
+    result = np.full((side, side, C), color, dtype=np.uint8)
+    result[(side-cur_h)//2:(side-cur_h)//2+cur_h, (side-cur_w)//2:(side - cur_w)//2+cur_w,:] = result_tmp
+    result = Image.fromarray(result)
+    return result.resize((size, size), Image.LANCZOS) if size else result

mvd/hunyuan3d_mvd_lite_pipeline.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 
@@ -62,10 +64,10 @@ EXAMPLE_DOC_STRING = """
     Examples:
         ```py
         >>> import torch
-        >>> from here import Hunyuan3d_MVD_Qing_Pipeline
+        >>> from here import Hunyuan3d_MVD_Lite_Pipeline
 
-        >>> pipe = Hunyuan3d_MVD_Qing_Pipeline.from_pretrained(
-        ...     "Tencent-Hunyuan-3D/MVD-Qing", torch_dtype=torch.float16
+        >>> pipe = Hunyuan3d_MVD_Lite_Pipeline.from_pretrained(
+        ...     "weights/mvd_lite", torch_dtype=torch.float16
         ... )
         >>> pipe.to("cuda")
 
@@ -173,18 +175,17 @@ class Hunyuan3d_MVD_Lite_Pipeline(DiffusionPipeline, TextualInversionLoaderMixin
             text_encoder=text_encoder,
             vision_encoder=vision_encoder,
             feature_extractor_vae=feature_extractor_vae,
-            feature_extractor_clip=feature_extractor_clip)
-        '''
-            rewrite the stable diffusion pipeline
-            vae: vae
-            unet: unet
-            tokenizer: tokenizer
-            scheduler: scheduler
-            text_encoder: text_encoder
-            vision_encoder: vision_encoder
-            feature_extractor_vae: feature_extractor_vae
-            feature_extractor_clip: feature_extractor_clip
-        '''
+            feature_extractor_clip=feature_extractor_clip
+        )
+            # rewrite the stable diffusion pipeline
+            # vae: vae
+            # unet: unet
+            # tokenizer: tokenizer
+            # scheduler: scheduler
+            # text_encoder: text_encoder
+            # vision_encoder: vision_encoder
+            # feature_extractor_vae: feature_extractor_vae
+            # feature_extractor_clip: feature_extractor_clip
         self.register_to_config(ramping_coefficients=ramping_coefficients)
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

mvd/hunyuan3d_mvd_std_pipeline.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

mvd/utils.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 

svrm/.ipynb_checkpoints/predictor-checkpoint.py

@@ -0,0 +1,152 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os
+import math
+import time
+import torch
+import numpy as np
+from tqdm import tqdm
+from PIL import Image, ImageSequence
+from omegaconf import OmegaConf
+from torchvision import transforms
+from safetensors.torch import save_file, load_file
+from .ldm.util import instantiate_from_config
+from .ldm.vis_util import render
+
+class MV23DPredictor(object):
+    def __init__(self, ckpt_path, cfg_path, elevation=15, number_view=60, 
+                render_size=256, device="cuda:0") -> None:
+        self.device = device
+        self.elevation = elevation
+        self.number_view = number_view
+        self.render_size = render_size
+
+        self.elevation_list = [0, 0, 0, 0, 0, 0, 0]
+        self.azimuth_list = [0, 60, 120, 180, 240, 300, 0]
+
+        st = time.time()
+        self.model = self.init_model(ckpt_path, cfg_path)
+        print(f"=====> mv23d model init time: {time.time() - st}")
+
+        self.input_view_transform = transforms.Compose([
+            transforms.Resize(504, interpolation=Image.BICUBIC),
+            transforms.ToTensor(),
+        ])
+        self.final_input_view_transform = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+
+    def init_model(self, ckpt_path, cfg_path):
+        config = OmegaConf.load(cfg_path)
+        model = instantiate_from_config(config.model)
+
+        weights = load_file("./weights/svrm/svrm.safetensors")
+        model.load_state_dict(weights)
+
+        model.to(self.device)
+        model = model.eval()
+        model.render.half()
+        print(f'Load model successfully')
+        return model
+
+    def create_camera_to_world_matrix(self, elevation, azimuth, cam_dis=1.5):
+        # elevation azimuth are radians
+        # Convert elevation and azimuth angles to Cartesian coordinates on a unit sphere
+        x = np.cos(elevation) * np.cos(azimuth)
+        y = np.cos(elevation) * np.sin(azimuth)
+        z = np.sin(elevation)
+
+        # Calculate camera position, target, and up vectors
+        camera_pos = np.array([x, y, z]) * cam_dis
+        target = np.array([0, 0, 0])
+        up = np.array([0, 0, 1])
+
+        # Construct view matrix
+        forward = target - camera_pos
+        forward /= np.linalg.norm(forward)
+        right = np.cross(forward, up)
+        right /= np.linalg.norm(right)
+        new_up = np.cross(right, forward)
+        new_up /= np.linalg.norm(new_up)
+        cam2world = np.eye(4)
+        cam2world[:3, :3] = np.array([right, new_up, -forward]).T
+        cam2world[:3, 3] = camera_pos
+        return cam2world
+
+    def refine_mask(self, mask, k=16):
+        mask /= 255.0
+        boder_mask = (mask >= -math.pi / 2.0 / k + 0.5) & (mask <= math.pi / 2.0 / k + 0.5)
+        mask[boder_mask] = 0.5 * np.sin(k * (mask[boder_mask] - 0.5)) + 0.5
+        mask[mask < -math.pi / 2.0 / k + 0.5] = 0.0
+        mask[mask > math.pi / 2.0 / k + 0.5] = 1.0
+        return (mask * 255.0).astype(np.uint8)
+
+    def load_images_and_cameras(self, input_imgs, elevation_list, azimuth_list):
+        input_image_list = []
+        input_cam_list = []
+        for input_view_image, elevation, azimuth in zip(input_imgs, elevation_list, azimuth_list):
+            input_view_image = self.input_view_transform(input_view_image)  
+            input_image_list.append(input_view_image)
+
+            input_view_cam_pos = self.create_camera_to_world_matrix(np.radians(elevation), np.radians(azimuth))
+            input_view_cam_intrinsic = np.array([35. / 32, 35. /32, 0.5, 0.5])
+            input_view_cam = torch.from_numpy(
+                np.concatenate([input_view_cam_pos.reshape(-1), input_view_cam_intrinsic], 0)
+            ).float()
+            input_cam_list.append(input_view_cam)
+
+        pixels_input = torch.stack(input_image_list, dim=0)
+        input_images = self.final_input_view_transform(pixels_input)
+        input_cams = torch.stack(input_cam_list, dim=0)
+        return input_images, input_cams
+
+    def load_data(self, intput_imgs):
+        assert (6+1) == len(intput_imgs)
+        
+        input_images, input_cams = self.load_images_and_cameras(intput_imgs, self.elevation_list, self.azimuth_list)
+        input_cams[-1, :] = 0 # for user input view
+        
+        data = {}
+        data["input_view"] = input_images.unsqueeze(0).to(self.device)    # 1 4 3 512 512
+        data["input_view_cam"] = input_cams.unsqueeze(0).to(self.device)  # 1 4 20
+        return data
+
+    @torch.no_grad()
+    def predict(
+        self, 
+        intput_imgs, 
+        save_dir = "outputs/", 
+        image_input = None,
+        target_face_count = 10000,
+        do_texture_mapping = True,
+    ):
+        os.makedirs(save_dir, exist_ok=True)
+        print(save_dir)
+        
+        with torch.cuda.amp.autocast():
+            self.model.export_mesh_with_uv(
+                data = self.load_data(intput_imgs),
+                out_dir = save_dir,
+                target_face_count = target_face_count,
+                do_texture_mapping = do_texture_mapping
+            )

svrm/ldm/.ipynb_checkpoints/util-checkpoint.py

@@ -0,0 +1,252 @@
+import os
+import importlib
+from inspect import isfunction
+import cv2
+import time
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import matplotlib.pyplot as plt
+import torch
+from torch import optim
+import torchvision
+
+
+def pil_rectangle_crop(im):
+    width, height = im.size   # Get dimensions
+    
+    if width <= height:
+        left = 0
+        right = width
+        top = (height - width)/2
+        bottom = (height + width)/2
+    else:
+        
+        top = 0
+        bottom = height
+        left = (width - height) / 2
+        bottom = (width + height) / 2
+
+    # Crop the center of the image
+    im = im.crop((left, top, right, bottom))
+    return im
+
+
+def add_margin(pil_img, color, size=256):
+    width, height = pil_img.size
+    result = Image.new(pil_img.mode, (size, size), color)
+    result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
+    return result
+
+
+def load_and_preprocess(interface, input_im):
+    '''
+    :param input_im (PIL Image).
+    :return image (H, W, 3) array in [0, 1].
+    '''
+    # See https://github.com/Ir1d/image-background-remove-tool
+    image = input_im.convert('RGB')
+
+    image_without_background = interface([image])[0]
+    image_without_background = np.array(image_without_background)
+    est_seg = image_without_background > 127
+    image = np.array(image)
+    foreground = est_seg[:, : , -1].astype(np.bool_)
+    image[~foreground] = [255., 255., 255.]
+    x, y, w, h = cv2.boundingRect(foreground.astype(np.uint8))
+    image = image[y:y+h, x:x+w, :]
+    image = Image.fromarray(np.array(image))
+    
+    # resize image such that long edge is 512
+    image.thumbnail([200, 200], Image.Resampling.LANCZOS)
+    image = add_margin(image, (255, 255, 255), size=256)
+    image = np.array(image)
+    return image
+
+
+def log_txt_as_img(wh, xc, size=10):
+    # wh a tuple of (width, height)
+    # xc a list of captions to plot
+    b = len(xc)
+    txts = list()
+    for bi in range(b):
+        txt = Image.new("RGB", wh, color="white")
+        draw = ImageDraw.Draw(txt)
+        font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
+        nc = int(40 * (wh[0] / 256))
+        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
+
+        try:
+            draw.text((0, 0), lines, fill="black", font=font)
+        except UnicodeEncodeError:
+            print("Cant encode string for logging. Skipping.")
+
+        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+        txts.append(txt)
+    txts = np.stack(txts)
+    txts = torch.tensor(txts)
+    return txts
+
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+
+def isimage(x):
+    if not isinstance(x,torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def mean_flat(tensor):
+    """
+    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def count_params(model, verbose=False):
+    total_params = sum(p.numel() for p in model.parameters())
+    if verbose:
+        print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
+    return total_params
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == '__is_first_stage__':
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+class AdamWwithEMAandWings(optim.Optimizer):
+    # credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
+    def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8,  # TODO: check hyperparameters before using
+                 weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999,   # ema decay to match previous code
+                 ema_power=1., param_names=()):
+        """AdamW that saves EMA versions of the parameters."""
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= ema_decay <= 1.0:
+            raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
+                        ema_power=ema_power, param_names=param_names)
+        super().__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            ema_params_with_grad = []
+            state_sums = []
+            max_exp_avg_sqs = []
+            state_steps = []
+            amsgrad = group['amsgrad']
+            beta1, beta2 = group['betas']
+            ema_decay = group['ema_decay']
+            ema_power = group['ema_power']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                params_with_grad.append(p)
+                if p.grad.is_sparse:
+                    raise RuntimeError('AdamW does not support sparse gradients')
+                grads.append(p.grad)
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of parameter values
+                    state['param_exp_avg'] = p.detach().float().clone()
+
+                exp_avgs.append(state['exp_avg'])
+                exp_avg_sqs.append(state['exp_avg_sq'])
+                ema_params_with_grad.append(state['param_exp_avg'])
+
+                if amsgrad:
+                    max_exp_avg_sqs.append(state['max_exp_avg_sq'])
+
+                # update the steps for each param group update
+                state['step'] += 1
+                # record the step after step update
+                state_steps.append(state['step'])
+
+            optim._functional.adamw(params_with_grad,
+                    grads,
+                    exp_avgs,
+                    exp_avg_sqs,
+                    max_exp_avg_sqs,
+                    state_steps,
+                    amsgrad=amsgrad,
+                    beta1=beta1,
+                    beta2=beta2,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    eps=group['eps'],
+                    maximize=False)
+
+            cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
+            for param, ema_param in zip(params_with_grad, ema_params_with_grad):
+                ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
+
+        return loss

svrm/ldm/models/.ipynb_checkpoints/svrm-checkpoint.py

@@ -0,0 +1,281 @@
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
+# The below software and/or models in this distribution may have been 
+# modified by THL A29 Limited ("Tencent Modifications"). 
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
+# except for the third-party components listed below. 
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
+# in the repsective licenses of these third-party components. 
+# Users must comply with all terms and conditions of original licenses of these third-party 
+# components and must ensure that the usage of the third party components adheres to 
+# all relevant laws and regulations. 
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and 
+# their software and algorithms, including trained model weights, parameters (including 
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
+# fine-tuning enabling code and other elements of the foregoing made publicly available 
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os
+import time
+import math
+import cv2
+import numpy as np
+import itertools
+import shutil
+from tqdm import tqdm
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+try:
+    import trimesh
+    import mcubes
+    import xatlas
+    import open3d as o3d
+except:
+    raise "failed to import 3d libraries "
+
+from ..modules.rendering_neus.mesh import Mesh
+from ..modules.rendering_neus.rasterize import NVDiffRasterizerContext
+
+from ..utils.ops import scale_tensor
+from ..util import count_params, instantiate_from_config
+from ..vis_util import render
+
+
+def unwrap_uv(v_pos, t_pos_idx):
+    print("Using xatlas to perform UV unwrapping, may take a while ...")
+    atlas = xatlas.Atlas()
+    atlas.add_mesh(v_pos, t_pos_idx)
+    atlas.generate(xatlas.ChartOptions(), xatlas.PackOptions())
+    _, indices, uvs = atlas.get_mesh(0)
+    indices = indices.astype(np.int64, casting="same_kind")
+    return uvs, indices
+
+
+def uv_padding(image, hole_mask, uv_padding_size = 2):
+    return cv2.inpaint(
+        (image.detach().cpu().numpy() * 255).astype(np.uint8),
+        (hole_mask.detach().cpu().numpy() * 255).astype(np.uint8),
+        uv_padding_size,
+        cv2.INPAINT_TELEA
+    )
+
+def refine_mesh(vtx_refine, faces_refine):
+    mesh = o3d.geometry.TriangleMesh(
+        vertices=o3d.utility.Vector3dVector(vtx_refine), 
+        triangles=o3d.utility.Vector3iVector(faces_refine)
+    )
+
+    mesh = mesh.remove_unreferenced_vertices()
+    mesh = mesh.remove_duplicated_triangles()
+    mesh = mesh.remove_duplicated_vertices()
+
+    voxel_size = max(mesh.get_max_bound() - mesh.get_min_bound())
+
+    mesh = mesh.simplify_vertex_clustering(
+        voxel_size=0.007, # 0.005
+        contraction=o3d.geometry.SimplificationContraction.Average)
+
+    mesh = mesh.filter_smooth_simple(number_of_iterations=2)
+
+    vtx_refine = np.asarray(mesh.vertices).astype(np.float32)
+    faces_refine = np.asarray(mesh.triangles)
+    return vtx_refine, faces_refine, mesh
+
+
+class SVRMModel(torch.nn.Module):
+    def __init__(
+        self,
+        img_encoder_config,
+        img_to_triplane_config,
+        render_config,
+        device = "cuda:0",
+        **kwargs
+    ):
+        super(SVRMModel, self).__init__()
+        self.img_encoder = instantiate_from_config(img_encoder_config).half()
+        self.img_to_triplane_decoder = instantiate_from_config(img_to_triplane_config).half()
+        self.render = instantiate_from_config(render_config).half()
+        self.device = device
+        count_params(self, verbose=True)
+        
+
+    @torch.no_grad()
+    def export_mesh_with_uv(
+        self, 
+        data, 
+        mesh_size: int = 384, 
+        ctx = None, 
+        context_type = 'cuda', 
+        texture_res = 1024,
+        target_face_count = 10000,
+        do_texture_mapping = True,
+        out_dir = 'outputs/test'
+    ):
+        """
+        color_type: 0 for ray texture, 1 for vertices texture
+        """
+        
+        obj_vertext_path = os.path.join(out_dir, 'mesh_with_colors.obj')
+        obj_path = os.path.join(out_dir, 'mesh.obj')
+        obj_texture_path = os.path.join(out_dir, 'texture.png')
+        obj_mtl_path = os.path.join(out_dir, 'texture.mtl')
+        glb_path = os.path.join(out_dir, 'mesh.glb')
+
+        st = time.time()
+        
+        here = {'device': self.device, 'dtype': torch.float16}
+        input_view_image = data["input_view"].to(**here)    # [b, m, c, h, w]
+        input_view_cam = data["input_view_cam"].to(**here)  # [b, m, 20]
+
+        batch_size, input_view_num, *_ = input_view_image.shape
+        assert batch_size == 1, "batch size should be 1"
+
+        input_view_image = rearrange(input_view_image, 'b m c h w -> (b m) c h w')
+        input_view_cam = rearrange(input_view_cam, 'b m d -> (b m) d')
+        input_view_feat = self.img_encoder(input_view_image, input_view_cam)
+        input_view_feat = rearrange(input_view_feat, '(b m) l d -> b (l m) d', m=input_view_num)
+
+        # -- decoder
+        torch.cuda.empty_cache()
+        triplane_gen = self.img_to_triplane_decoder(input_view_feat)  # [b, 3, tri_dim, h, w]
+        del input_view_feat
+        torch.cuda.empty_cache()
+
+        # --- triplane nerf render
+
+        cur_triplane = triplane_gen[0:1]
+        
+        aabb = torch.tensor([[-0.6, -0.6, -0.6], [0.6, 0.6, 0.6]]).unsqueeze(0).to(**here)
+        grid_out = self.render.forward_grid(planes=cur_triplane, grid_size=mesh_size, aabb=aabb)
+
+        print(f"=====> Triplane forward time: {time.time() - st}")
+        st = time.time()
+        
+        vtx, faces = mcubes.marching_cubes(0. - grid_out['sdf'].squeeze(0).squeeze(-1).cpu().float().numpy(), 0)
+        
+        bbox = aabb[0].cpu().numpy()
+        vtx = vtx / (mesh_size - 1)  
+        vtx = vtx * (bbox[1] - bbox[0]) + bbox[0]
+
+        # refine mesh
+        vtx_refine, faces_refine, mesh = refine_mesh(vtx, faces)
+
+        # reduce faces
+        if faces_refine.shape[0] > target_face_count:
+            print(f"reduce face: {faces_refine.shape[0]} -> {target_face_count}")
+            mesh = o3d.geometry.TriangleMesh(
+                vertices = o3d.utility.Vector3dVector(vtx_refine),
+                triangles = o3d.utility.Vector3iVector(faces_refine)
+            )
+            
+            # Function to simplify mesh using Quadric Error Metric Decimation by Garland and Heckbert
+            mesh = mesh.simplify_quadric_decimation(target_face_count, boundary_weight=1.0)
+
+            mesh = Mesh(
+                v_pos = torch.from_numpy(np.asarray(mesh.vertices)).to(self.device),
+                t_pos_idx = torch.from_numpy(np.asarray(mesh.triangles)).to(self.device),
+                v_rgb = torch.from_numpy(np.asarray(mesh.vertex_colors)).to(self.device)
+            )
+            vtx_refine = mesh.v_pos.cpu().numpy()
+            faces_refine = mesh.t_pos_idx.cpu().numpy()
+
+        vtx_colors = self.render.forward_points(cur_triplane, torch.tensor(vtx_refine).unsqueeze(0).to(**here))
+        vtx_colors = vtx_colors['rgb'].float().squeeze(0).cpu().numpy()
+
+        color_ratio = 0.8 # increase brightness
+        with open(obj_vertext_path, 'w') as fid:
+            verts = vtx_refine[:, [1,2,0]] 
+            for pidx, pp in enumerate(verts):
+                color = vtx_colors[pidx]
+                color = [color[0]**color_ratio, color[1]**color_ratio, color[2]**color_ratio]
+                fid.write('v %f %f %f %f %f %f\n' % (pp[0], pp[1], pp[2], color[0], color[1], color[2]))
+            for i, f in enumerate(faces_refine):
+                f1 = f + 1
+                fid.write('f %d %d %d\n' % (f1[0], f1[1], f1[2]))
+                
+        mesh = trimesh.load_mesh(obj_vertext_path)
+        print(f"=====> generate mesh with vertex shading time: {time.time() - st}")
+        st = time.time()
+
+        if not do_texture_mapping:
+            shutil.copy(obj_vertext_path, obj_path)
+            mesh.export(glb_path, file_type='glb')
+            return None
+            
+
+        ##########  export texture  ########
+        
+        
+        st = time.time()
+        
+        # uv unwrap 
+        vtx_tex, t_tex_idx = unwrap_uv(vtx_refine, faces_refine)      
+        vtx_refine   = torch.from_numpy(vtx_refine).to(self.device)   
+        faces_refine = torch.from_numpy(faces_refine).to(self.device)  
+        t_tex_idx    = torch.from_numpy(t_tex_idx).to(self.device)    
+        uv_clip      = torch.from_numpy(vtx_tex * 2.0 - 1.0).to(self.device) 
+
+        # rasterize
+        ctx = NVDiffRasterizerContext(context_type, cur_triplane.device) if ctx is None else ctx
+        rast = ctx.rasterize_one(
+            torch.cat([
+                uv_clip, 
+                torch.zeros_like(uv_clip[..., 0:1]), 
+                torch.ones_like(uv_clip[..., 0:1])
+            ], dim=-1), 
+            t_tex_idx,  
+            (texture_res, texture_res)
+        )[0]
+        hole_mask = ~(rast[:, :, 3] > 0)
+
+        # Interpolate world space position
+        gb_pos = ctx.interpolate_one(vtx_refine, rast[None, ...], faces_refine)[0][0]
+        
+        with torch.no_grad():
+            gb_mask_pos_scale = scale_tensor(gb_pos.unsqueeze(0).view(1, -1, 3), (-1, 1), (-1, 1))
+            
+            tex_map = self.render.forward_points(cur_triplane, gb_mask_pos_scale)['rgb']
+            
+            tex_map = tex_map.float().squeeze(0)  # (0, 1)
+            tex_map = tex_map.view((texture_res, texture_res, 3)) 
+            img = uv_padding(tex_map, hole_mask)
+            img = ((img/255.0) ** color_ratio) * 255  # increase brightness
+            img = img.clip(0, 255).astype(np.uint8)
+    
+        verts = vtx_refine.cpu().numpy()[:, [1,2,0]] 
+        faces = faces_refine.cpu().numpy()
+
+        with open(obj_mtl_path, 'w') as fid:
+            fid.write('newmtl material_0\n')
+            fid.write("Ka 1.000 1.000 1.000\n")
+            fid.write("Kd 1.000 1.000 1.000\n")
+            fid.write("Ks 0.000 0.000 0.000\n")
+            fid.write("d 1.0\n")
+            fid.write("illum 2\n")
+            fid.write(f'map_Kd texture.png\n')
+        
+        with open(obj_path, 'w') as fid:
+            fid.write(f'mtllib texture.mtl\n')
+            for pidx, pp in enumerate(verts):
+                fid.write('v %f %f %f\n' % (pp[0], pp[1], pp[2]))
+            for pidx, pp in enumerate(vtx_tex): 
+                fid.write('vt %f %f\n' % (pp[0], 1 - pp[1]))
+            fid.write('usemtl material_0\n')
+            for i, f in enumerate(faces):
+                f1 = f + 1
+                f2 = t_tex_idx[i] + 1
+                fid.write('f %d/%d %d/%d %d/%d\n' % (f1[0], f2[0], f1[1], f2[1], f1[2], f2[2],))
+
+        cv2.imwrite(obj_texture_path, img[..., [2, 1, 0]])   
+        mesh = trimesh.load_mesh(obj_path)
+        mesh.export(glb_path, file_type='glb')
+        print(f"=====> generate mesh with texture shading time: {time.time() - st}")
+  

svrm/ldm/models/svrm.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been 
@@ -68,7 +70,8 @@ def uv_padding(image, hole_mask, uv_padding_size = 2):
 def refine_mesh(vtx_refine, faces_refine):
     mesh = o3d.geometry.TriangleMesh(
         vertices=o3d.utility.Vector3dVector(vtx_refine), 
-        triangles=o3d.utility.Vector3iVector(faces_refine))
+        triangles=o3d.utility.Vector3iVector(faces_refine)
+    )
 
     mesh = mesh.remove_unreferenced_vertices()
     mesh = mesh.remove_duplicated_triangles()
@@ -235,9 +238,12 @@ class SVRMModel(torch.nn.Module):
 
         # Interpolate world space position
         gb_pos = ctx.interpolate_one(vtx_refine, rast[None, ...], faces_refine)[0][0]
+        
         with torch.no_grad():
             gb_mask_pos_scale = scale_tensor(gb_pos.unsqueeze(0).view(1, -1, 3), (-1, 1), (-1, 1))
+            
             tex_map = self.render.forward_points(cur_triplane, gb_mask_pos_scale)['rgb']
+            
             tex_map = tex_map.float().squeeze(0)  # (0, 1)
             tex_map = tex_map.view((texture_res, texture_res, 3)) 
             img = uv_padding(tex_map, hole_mask)

svrm/predictor.py

@@ -1,5 +1,7 @@
-# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
-# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+# Open Source Model Licensed under the Apache License Version 2.0 
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited 
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
 
 # Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
 # The below software and/or models in this distribution may have been