Update app.py

app.py

@@ -60,7 +60,91 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     processed_image.save(f"{TMP_DIR}/{trial_id}.png")
     return trial_id, processed_image
 
-[이전 코드의 나머지 함수들: pack_state, unpack_state, image_to_3d, extract_glb, activate_button, deactivate_button은 그대로 유지]
+def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
+    return {
+        'gaussian': {
+            **gs.init_params,
+            '_xyz': gs._xyz.cpu().numpy(),
+            '_features_dc': gs._features_dc.cpu().numpy(),
+            '_scaling': gs._scaling.cpu().numpy(),
+            '_rotation': gs._rotation.cpu().numpy(),
+            '_opacity': gs._opacity.cpu().numpy(),
+        },
+        'mesh': {
+            'vertices': mesh.vertices.cpu().numpy(),
+            'faces': mesh.faces.cpu().numpy(),
+        },
+        'trial_id': trial_id,
+    }
+    
+    
+def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
+    gs = Gaussian(
+        aabb=state['gaussian']['aabb'],
+        sh_degree=state['gaussian']['sh_degree'],
+        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+        scaling_bias=state['gaussian']['scaling_bias'],
+        opacity_bias=state['gaussian']['opacity_bias'],
+        scaling_activation=state['gaussian']['scaling_activation'],
+    )
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+    
+    mesh = edict(
+        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    )
+    
+    return gs, mesh, state['trial_id']
+
+
+@spaces.GPU
+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]:
+    if randomize_seed:
+        seed = np.random.randint(0, MAX_SEED)
+    outputs = pipeline.run(
+        Image.open(f"{TMP_DIR}/{trial_id}.png"),
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+    )
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    trial_id = uuid.uuid4()
+    video_path = f"{TMP_DIR}/{trial_id}.mp4"
+    os.makedirs(os.path.dirname(video_path), exist_ok=True)
+    imageio.mimsave(video_path, video, fps=15)
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
+    return state, video_path
+
+
+@spaces.GPU
+def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]:
+    gs, mesh, trial_id = unpack_state(state)
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+    glb_path = f"{TMP_DIR}/{trial_id}.glb"
+    glb.export(glb_path)
+    return glb_path, glb_path
+
+
+def activate_button() -> gr.Button:
+    return gr.Button(interactive=True)
+
+
+def deactivate_button() -> gr.Button:
+    return gr.Button(interactive=False)
 
 @spaces.GPU
 def text_to_image(prompt: str, height: int, width: int, steps: int, scales: float, seed: int) -> Image.Image: