File size: 2,753 Bytes
2e90551
 
 
 
921d811
2e90551
 
 
 
 
 
 
 
 
 
 
 
 
 
9b53de6
380617c
2e90551
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9dd1694
 
 
 
 
 
 
 
2e90551
9dd1694
 
 
 
 
 
2e90551
 
 
 
 
 
 
 
 
 
 
 
 
 
5aae6b8
2e90551
044ab04
 
 
 
 
 
 
 
 
 
 
 
 
 
2e90551
044ab04
 
2e90551
044ab04
 
 
2e90551
 
 
 
921d811
cac4e12
2e90551
ad80ffb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import os
import shutil
from functools import partial

import gradio as gr
from common import (
    MAX_SEED,
    VERSION,
    TrellisImageTo3DPipeline,
    active_btn_by_content,
    extract_3d_representations_v2,
    extract_urdf,
    get_seed,
    image_to_3d,
    preprocess_image_fn,
    preprocess_sam_image_fn,
    select_point,
)
from gradio.themes import Default
from gradio.themes.utils.colors import slate
from gradio_litmodel3d import LitModel3D
from asset3d_gen.models.delight import DelightingModel
from asset3d_gen.models.segment import RembgRemover, SAMPredictor
from asset3d_gen.models.super_resolution import ImageRealESRGAN
from asset3d_gen.utils.gpt_clients import GPT_CLIENT
from asset3d_gen.validators.quality_checkers import (
    ImageAestheticChecker,
    ImageSegChecker,
    MeshGeoChecker,
)
from asset3d_gen.validators.urdf_convertor import URDFGenerator

TMP_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
)
os.makedirs(TMP_DIR, exist_ok=True)

# RBG_REMOVER = RembgRemover()
# SAM_PREDICTOR = SAMPredictor(model_type="vit_h")
# DELIGHT = DelightingModel()
# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
# PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
#     "JeffreyXiang/TRELLIS-image-large"
# )
# # PIPELINE.cuda()

# IMAGE_BUFFER = {}
# SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
# GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
# AESTHETIC_CHECKER = ImageAestheticChecker()
# CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
# URDF_CONVERTOR = URDFGenerator(GPT_CLIENT, render_view_num=4)


def start_session(req: gr.Request) -> None:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)


def end_session(req: gr.Request) -> None:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    if os.path.exists(user_dir):
        shutil.rmtree(user_dir)


with gr.Blocks(
    delete_cache=(43200, 43200)
) as demo:
    with gr.Column(scale=1):
        video_output = gr.Video(
            label="Generated 3D Asset",
            autoplay=True,
            loop=True,
            height=300,
        )
        model_output_gs = gr.Model3D(
            label="Gaussian Representation", height=300, interactive=False
        )
        aligned_gs = gr.Textbox(visible=False)
        with gr.Row():
            model_output_mesh = LitModel3D(
                label="Mesh Representation",
                height=300,
                exposure=10,
                interactive=False
            )
        gr.Markdown(
            """ The rendering of `Gaussian Representation` takes additional 10s. """  # noqa
        )

    demo.load(start_session)
    demo.unload(end_session)



if __name__ == "__main__":
    demo.launch()