yejunliang23 commited on
Commit
981f763
·
verified ·
1 Parent(s): c950af5

Delete app1.py

Browse files
Files changed (1) hide show
  1. app1.py +0 -264
app1.py DELETED
@@ -1,264 +0,0 @@
1
- import gradio as gr
2
- import spaces
3
- from gradio_litmodel3d import LitModel3D
4
-
5
- import os
6
- os.environ['SPCONV_ALGO'] = 'native'
7
- from typing import *
8
- import torch
9
- import numpy as np
10
- import imageio
11
- import uuid
12
- from easydict import EasyDict as edict
13
- from PIL import Image
14
- from trellis.pipelines import TrellisImageTo3DPipeline
15
- from trellis.representations import Gaussian, MeshExtractResult
16
- from trellis.utils import render_utils, postprocessing_utils
17
-
18
- import logging
19
-
20
- # Configure logging
21
- logging.basicConfig(
22
- level=logging.INFO,
23
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
24
- handlers=[
25
- logging.StreamHandler()
26
- ]
27
- )
28
- logger = logging.getLogger(__name__)
29
-
30
- # Log environment variables
31
- logger.info(f"ATTN_BACKEND: {os.environ.get('ATTN_BACKEND')}")
32
- logger.info(f"ATTN_DEBUG: {os.environ.get('ATTN_DEBUG')}")
33
- logger.info(f"SPARSE_BACKEND: {os.environ.get('SPARSE_BACKEND')}")
34
- logger.info(f"SPARSE_DEBUG: {os.environ.get('SPARSE_DEBUG')}")
35
- logger.info(f"SPARSE_ATTN_BACKEND: {os.environ.get('SPARSE_ATTN_BACKEND')}")
36
-
37
- MAX_SEED = np.iinfo(np.int32).max
38
- TMP_DIR = "/tmp/Trellis-demo"
39
-
40
- os.makedirs(TMP_DIR, exist_ok=True)
41
-
42
-
43
- def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
44
- """
45
- Preprocess the input image.
46
- Args:
47
- image (Image.Image): The input image.
48
- Returns:
49
- str: uuid of the trial.
50
- Image.Image: The preprocessed image.
51
- """
52
- trial_id = str(uuid.uuid4())
53
- processed_image = pipeline.preprocess_image(image)
54
- processed_image.save(f"{TMP_DIR}/{trial_id}.png")
55
- return trial_id, processed_image
56
-
57
-
58
- def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
59
- return {
60
- 'gaussian': {
61
- **gs.init_params,
62
- '_xyz': gs._xyz.cpu().numpy(),
63
- '_features_dc': gs._features_dc.cpu().numpy(),
64
- '_scaling': gs._scaling.cpu().numpy(),
65
- '_rotation': gs._rotation.cpu().numpy(),
66
- '_opacity': gs._opacity.cpu().numpy(),
67
- },
68
- 'mesh': {
69
- 'vertices': mesh.vertices.cpu().numpy(),
70
- 'faces': mesh.faces.cpu().numpy(),
71
- },
72
- 'trial_id': trial_id,
73
- }
74
-
75
-
76
- def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
77
- gs = Gaussian(
78
- aabb=state['gaussian']['aabb'],
79
- sh_degree=state['gaussian']['sh_degree'],
80
- mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
81
- scaling_bias=state['gaussian']['scaling_bias'],
82
- opacity_bias=state['gaussian']['opacity_bias'],
83
- scaling_activation=state['gaussian']['scaling_activation'],
84
- )
85
- gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
86
- gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
87
- gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
88
- gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
89
- gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
90
-
91
- mesh = edict(
92
- vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
93
- faces=torch.tensor(state['mesh']['faces'], device='cuda'),
94
- )
95
-
96
- return gs, mesh, state['trial_id']
97
-
98
-
99
- @spaces.GPU
100
- def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
101
- """
102
- Convert an image to a 3D model.
103
- Args:
104
- trial_id (str): The uuid of the trial.
105
- seed (int): The random seed.
106
- randomize_seed (bool): Whether to randomize the seed.
107
- ss_guidance_strength (float): The guidance strength for sparse structure generation.
108
- ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
109
- slat_guidance_strength (float): The guidance strength for structured latent generation.
110
- slat_sampling_steps (int): The number of sampling steps for structured latent generation.
111
- Returns:
112
- dict: The information of the generated 3D model.
113
- str: The path to the video of the 3D model.
114
- """
115
- if randomize_seed:
116
- seed = np.random.randint(0, MAX_SEED)
117
- outputs = pipeline.run(
118
- Image.open(f"{TMP_DIR}/{trial_id}.png"),
119
- seed=seed,
120
- formats=["gaussian", "mesh"],
121
- preprocess_image=False,
122
- sparse_structure_sampler_params={
123
- "steps": ss_sampling_steps,
124
- "cfg_strength": ss_guidance_strength,
125
- },
126
- slat_sampler_params={
127
- "steps": slat_sampling_steps,
128
- "cfg_strength": slat_guidance_strength,
129
- },
130
- )
131
- video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
132
- video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
133
- video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
134
- trial_id = uuid.uuid4()
135
- video_path = f"{TMP_DIR}/{trial_id}.mp4"
136
- os.makedirs(os.path.dirname(video_path), exist_ok=True)
137
- imageio.mimsave(video_path, video, fps=15)
138
- state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
139
- return state, video_path
140
-
141
-
142
- @spaces.GPU
143
- def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]:
144
- """
145
- Extract a GLB file from the 3D model.
146
- Args:
147
- state (dict): The state of the generated 3D model.
148
- mesh_simplify (float): The mesh simplification factor.
149
- texture_size (int): The texture resolution.
150
- Returns:
151
- str: The path to the extracted GLB file.
152
- """
153
- gs, mesh, trial_id = unpack_state(state)
154
- glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
155
- glb_path = f"{TMP_DIR}/{trial_id}.glb"
156
- glb.export(glb_path)
157
- return glb_path, glb_path
158
-
159
-
160
- def activate_button() -> gr.Button:
161
- return gr.Button(interactive=True)
162
-
163
-
164
- def deactivate_button() -> gr.Button:
165
- return gr.Button(interactive=False)
166
-
167
-
168
- with gr.Blocks() as demo:
169
- gr.Markdown("""
170
- ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
171
- * Upload an image and click "Generate" to create a 3D asset. If the image has alpha channel, it be used as the mask. Otherwise, we use `rembg` to remove the background.
172
- * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it.
173
- """)
174
-
175
- with gr.Row():
176
- with gr.Column():
177
- image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil", height=300)
178
-
179
- with gr.Accordion(label="Generation Settings", open=False):
180
- seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
181
- randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
182
- gr.Markdown("Stage 1: Sparse Structure Generation")
183
- with gr.Row():
184
- ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
185
- ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
186
- gr.Markdown("Stage 2: Structured Latent Generation")
187
- with gr.Row():
188
- slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
189
- slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
190
-
191
- generate_btn = gr.Button("Generate")
192
-
193
- with gr.Accordion(label="GLB Extraction Settings", open=False):
194
- mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
195
- texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
196
-
197
- extract_glb_btn = gr.Button("Extract GLB", interactive=False)
198
-
199
- with gr.Column():
200
- video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
201
- model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300)
202
- download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
203
-
204
- trial_id = gr.Textbox(visible=False)
205
- output_buf = gr.State()
206
-
207
- # Example images at the bottom of the page
208
-
209
- # Handlers
210
- image_prompt.upload(
211
- preprocess_image,
212
- inputs=[image_prompt],
213
- outputs=[trial_id, image_prompt],
214
- )
215
- image_prompt.clear(
216
- lambda: '',
217
- outputs=[trial_id],
218
- )
219
-
220
- generate_btn.click(
221
- image_to_3d,
222
- inputs=[trial_id, seed, randomize_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
223
- outputs=[output_buf, video_output],
224
- ).then(
225
- activate_button,
226
- outputs=[extract_glb_btn],
227
- )
228
-
229
- video_output.clear(
230
- deactivate_button,
231
- outputs=[extract_glb_btn],
232
- )
233
-
234
- extract_glb_btn.click(
235
- extract_glb,
236
- inputs=[output_buf, mesh_simplify, texture_size],
237
- outputs=[model_output, download_glb],
238
- ).then(
239
- activate_button,
240
- outputs=[download_glb],
241
- )
242
-
243
- model_output.clear(
244
- deactivate_button,
245
- outputs=[download_glb],
246
- )
247
-
248
-
249
- # Launch the Gradio app
250
- if __name__ == "__main__":
251
- pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
252
- if torch.cuda.is_available():
253
- pipeline.cuda()
254
- print("CUDA is available. Using GPU.")
255
- else:
256
- print("CUDA not available. Falling back to CPU.")
257
- try:
258
- pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) # Preload rembg
259
- except:
260
- pass
261
- print(f"CUDA Available: {torch.cuda.is_available()}")
262
- print(f"CUDA Version: {torch.version.cuda}")
263
- print(f"Number of GPUs: {torch.cuda.device_count()}")
264
- demo.launch(debug=True)