cavargas10 commited on
Commit
0c561f4
verified
1 Parent(s): 701c397
Files changed (1) hide show
  1. app.py +46 -12
app.py CHANGED
@@ -12,22 +12,40 @@ from PIL import Image
12
  from trellis.pipelines import TrellisImageTo3DPipeline
13
  from trellis.representations import Gaussian, MeshExtractResult
14
  from trellis.utils import render_utils, postprocessing_utils
 
 
 
 
 
15
  MAX_SEED = np.iinfo(np.int32).max
16
  TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
17
  os.makedirs(TMP_DIR, exist_ok=True)
18
 
19
- # Funciones auxiliares
20
  def start_session(req: gr.Request):
21
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
 
 
22
  os.makedirs(user_dir, exist_ok=True)
23
 
24
  def end_session(req: gr.Request):
25
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
26
- shutil.rmtree(user_dir)
 
 
 
 
 
 
 
 
 
27
 
28
- def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
 
 
29
  images = [image[0] for image in images]
30
  processed_images = [pipeline.preprocess_image(image) for image in images]
 
31
  return processed_images
32
 
33
  def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
@@ -67,7 +85,9 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
67
  return gs, mesh
68
 
69
  def get_seed(randomize_seed: bool, seed: int) -> int:
70
- return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
 
71
 
72
  @spaces.GPU
73
  def image_to_3d(
@@ -80,12 +100,15 @@ def image_to_3d(
80
  multiimage_algo: Literal["multidiffusion", "stochastic"],
81
  req: gr.Request,
82
  ) -> Tuple[dict, str]:
83
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
 
 
 
84
  outputs = pipeline.run_multi_image(
85
  [image[0] for image in multiimages],
86
  seed=seed,
87
  formats=["gaussian", "mesh"],
88
- preprocess_image=False,
89
  sparse_structure_sampler_params={
90
  "steps": ss_sampling_steps,
91
  "cfg_strength": ss_guidance_strength,
@@ -96,13 +119,17 @@ def image_to_3d(
96
  },
97
  mode=multiimage_algo,
98
  )
 
 
99
  video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
100
  video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
101
  video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
102
  video_path = os.path.join(user_dir, 'sample.mp4')
103
  imageio.mimsave(video_path, video, fps=15)
 
104
  state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
105
  torch.cuda.empty_cache()
 
106
  return state, video_path
107
 
108
  @spaces.GPU(duration=90)
@@ -112,15 +139,19 @@ def extract_glb(
112
  texture_size: int,
113
  req: gr.Request,
114
  ) -> Tuple[str, str]:
115
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
 
 
 
116
  gs, mesh = unpack_state(state)
117
  glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
118
  glb_path = os.path.join(user_dir, 'sample.glb')
119
  glb.export(glb_path)
 
120
  torch.cuda.empty_cache()
 
121
  return glb_path, glb_path
122
 
123
- # Interfaz Gradio
124
  with gr.Blocks(delete_cache=(600, 600)) as demo:
125
  gr.Markdown("""
126
  # UTPL - Conversi贸n de Multiples Im谩genes a objetos 3D usando IA
@@ -164,24 +195,26 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
164
 
165
  output_buf = gr.State()
166
 
167
- # Manejadores
168
- demo.load(start_session)
169
  demo.unload(end_session)
170
 
171
  multiimage_prompt.upload(
172
  preprocess_images,
173
  inputs=[multiimage_prompt],
174
  outputs=[multiimage_prompt],
 
175
  )
176
 
177
  generate_btn.click(
178
  get_seed,
179
  inputs=[randomize_seed, seed],
180
  outputs=[seed],
 
181
  ).then(
182
  image_to_3d,
183
  inputs=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
184
  outputs=[output_buf, video_output],
 
185
  ).then(
186
  lambda: gr.Button(interactive=True),
187
  outputs=[extract_glb_btn],
@@ -196,6 +229,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
196
  extract_glb,
197
  inputs=[output_buf, mesh_simplify, texture_size],
198
  outputs=[model_output, download_glb],
 
199
  ).then(
200
  lambda: gr.Button(interactive=True),
201
  outputs=[download_glb],
 
12
  from trellis.pipelines import TrellisImageTo3DPipeline
13
  from trellis.representations import Gaussian, MeshExtractResult
14
  from trellis.utils import render_utils, postprocessing_utils
15
+
16
+ import logging
17
+
18
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - HF_SPACE_MULTIIMG - %(levelname)s - %(message)s')
19
+
20
  MAX_SEED = np.iinfo(np.int32).max
21
  TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
22
  os.makedirs(TMP_DIR, exist_ok=True)
23
 
 
24
  def start_session(req: gr.Request):
25
+ session_hash = str(req.session_hash)
26
+ user_dir = os.path.join(TMP_DIR, session_hash)
27
+ logging.info(f"START SESSION: Creando directorio para la sesi贸n {session_hash} en {user_dir}")
28
  os.makedirs(user_dir, exist_ok=True)
29
 
30
  def end_session(req: gr.Request):
31
+ session_hash = str(req.session_hash)
32
+ user_dir = os.path.join(TMP_DIR, session_hash)
33
+ logging.info(f"END SESSION: Intentando eliminar el directorio de la sesi贸n {session_hash} en {user_dir}")
34
+ if os.path.exists(user_dir):
35
+ try:
36
+ shutil.rmtree(user_dir)
37
+ logging.info(f"Directorio de la sesi贸n {session_hash} eliminado correctamente.")
38
+ except Exception as e:
39
+ logging.error(f"Error al eliminar el directorio de la sesi贸n {session_hash}: {e}")
40
+ else:
41
+ logging.warning(f"El directorio de la sesi贸n {session_hash} no fue encontrado al intentar eliminarlo. Es posible que ya haya sido limpiado.")
42
 
43
+ def preprocess_images(images: List[Tuple[Image.Image, str]], req: gr.Request) -> List[Image.Image]:
44
+ session_hash = str(req.session_hash)
45
+ logging.info(f"[{session_hash}] Preprocesando {len(images)} im谩genes.")
46
  images = [image[0] for image in images]
47
  processed_images = [pipeline.preprocess_image(image) for image in images]
48
+ logging.info(f"[{session_hash}] Preprocesamiento completado.")
49
  return processed_images
50
 
51
  def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
 
85
  return gs, mesh
86
 
87
  def get_seed(randomize_seed: bool, seed: int) -> int:
88
+ new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
89
+ logging.info(f"Usando seed: {new_seed}")
90
+ return new_seed
91
 
92
  @spaces.GPU
93
  def image_to_3d(
 
100
  multiimage_algo: Literal["multidiffusion", "stochastic"],
101
  req: gr.Request,
102
  ) -> Tuple[dict, str]:
103
+ session_hash = str(req.session_hash)
104
+ logging.info(f"[{session_hash}] Iniciando image_to_3d con {len(multiimages)} im谩genes.")
105
+ user_dir = os.path.join(TMP_DIR, session_hash)
106
+
107
  outputs = pipeline.run_multi_image(
108
  [image[0] for image in multiimages],
109
  seed=seed,
110
  formats=["gaussian", "mesh"],
111
+ preprocess_image=False,
112
  sparse_structure_sampler_params={
113
  "steps": ss_sampling_steps,
114
  "cfg_strength": ss_guidance_strength,
 
119
  },
120
  mode=multiimage_algo,
121
  )
122
+
123
+ logging.info(f"[{session_hash}] Generaci贸n del modelo completada. Renderizando video...")
124
  video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
125
  video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
126
  video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
127
  video_path = os.path.join(user_dir, 'sample.mp4')
128
  imageio.mimsave(video_path, video, fps=15)
129
+
130
  state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
131
  torch.cuda.empty_cache()
132
+ logging.info(f"[{session_hash}] Video renderizado y estado empaquetado. Devolviendo: {video_path}")
133
  return state, video_path
134
 
135
  @spaces.GPU(duration=90)
 
139
  texture_size: int,
140
  req: gr.Request,
141
  ) -> Tuple[str, str]:
142
+ session_hash = str(req.session_hash)
143
+ logging.info(f"[{session_hash}] Iniciando extract_glb...")
144
+ user_dir = os.path.join(TMP_DIR, session_hash)
145
+
146
  gs, mesh = unpack_state(state)
147
  glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
148
  glb_path = os.path.join(user_dir, 'sample.glb')
149
  glb.export(glb_path)
150
+
151
  torch.cuda.empty_cache()
152
+ logging.info(f"[{session_hash}] GLB extra铆do. Devolviendo: {glb_path}")
153
  return glb_path, glb_path
154
 
 
155
  with gr.Blocks(delete_cache=(600, 600)) as demo:
156
  gr.Markdown("""
157
  # UTPL - Conversi贸n de Multiples Im谩genes a objetos 3D usando IA
 
195
 
196
  output_buf = gr.State()
197
 
198
+ demo.load(start_session, inputs=None, outputs=None, api_name="start_session")
 
199
  demo.unload(end_session)
200
 
201
  multiimage_prompt.upload(
202
  preprocess_images,
203
  inputs=[multiimage_prompt],
204
  outputs=[multiimage_prompt],
205
+ api_name="preprocess_images"
206
  )
207
 
208
  generate_btn.click(
209
  get_seed,
210
  inputs=[randomize_seed, seed],
211
  outputs=[seed],
212
+ api_name="get_seed"
213
  ).then(
214
  image_to_3d,
215
  inputs=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
216
  outputs=[output_buf, video_output],
217
+ api_name="image_to_3d"
218
  ).then(
219
  lambda: gr.Button(interactive=True),
220
  outputs=[extract_glb_btn],
 
229
  extract_glb,
230
  inputs=[output_buf, mesh_simplify, texture_size],
231
  outputs=[model_output, download_glb],
232
+ api_name="extract_glb"
233
  ).then(
234
  lambda: gr.Button(interactive=True),
235
  outputs=[download_glb],