Update app.py

app.py

@@ -9,38 +9,57 @@ import logging
 import random
 import uuid
 
-# Configuración del logging para depuración
+# Configuración del logging para una mejor depuración
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')
 
+@spaces.GPU
 def install_dependencies():
-    """Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias."""
+    """
+    Instala de forma robusta el toolkit de CUDA y compila las extensiones C++/CUDA.
+    Usa subprocess.run para capturar errores.
+    """
     logging.info("Iniciando la instalación de dependencias...")
     
-    # Instalar CUDA Toolkit
+    # Instalar CUDA Toolkit si no está presente
     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run"
     CUDA_TOOLKIT_FILE = f"/tmp/{os.path.basename(CUDA_TOOLKIT_URL)}"
     if not os.path.exists("/usr/local/cuda"):
         logging.info("Descargando e instalando CUDA Toolkit...")
-        subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-        subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
-        subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+        subprocess.run(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE], check=True)
+        subprocess.run(["chmod", "+x", CUDA_TOOLKIT_FILE], check=True)
+        subprocess.run([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"], check=True)
     else:
         logging.info("CUDA Toolkit ya está instalado.")
 
+    # Configurar variables de entorno para la compilación
     os.environ["CUDA_HOME"] = "/usr/local/cuda"
-    os.environ["PATH"] = f"{os.environ['CUDA_HOME']}/bin:{os.environ['PATH']}"
-    os.environ["LD_LIBRARY_PATH"] = f"{os.environ['CUDA_HOME']}/lib:{os.environ.get('LD_LIBRARY_PATH', '')}"
+    os.environ["PATH"] = f"{os.environ.get('CUDA_HOME', '')}/bin:{os.environ.get('PATH', '')}"
+    os.environ["LD_LIBRARY_PATH"] = f"{os.environ.get('CUDA_HOME', '')}/lib:{os.environ.get('LD_LIBRARY_PATH', '')}"
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     
-    # Compilar extensiones personalizadas
-    logging.info("Compilando extensiones de renderizado...")
+    # Compilar extensiones personalizadas con manejo de errores
+    logging.info("Compilando extensión 'differentiable_renderer'...")
     renderer_path = "/home/user/app/step1x3d_texture/differentiable_renderer/"
-    subprocess.run(f"cd {renderer_path} && python setup.py install", shell=True, check=True)
-    subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
-    
-    logging.info("Instalación completada.")
-    os.system('nvcc -V')
-
+    try:
+        subprocess.run(f"cd {renderer_path} && python setup.py install", shell=True, check=True, capture_output=True, text=True)
+        logging.info("Extensión 'differentiable_renderer' compilada con éxito.")
+    except subprocess.CalledProcessError as e:
+        logging.error("¡FALLÓ LA COMPILACIÓN de 'differentiable_renderer'!")
+        logging.error(f"STDOUT: {e.stdout}")
+        logging.error(f"STDERR: {e.stderr}")
+        raise  # Detiene la aplicación si la compilación falla
+
+    try:
+        subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
+        logging.info("Extensión 'custom_rasterizer' instalada con éxito.")
+    except subprocess.CalledProcessError as e:
+        logging.error("¡FALLÓ LA INSTALACIÓN de 'custom_rasterizer'!")
+        raise
+
+    logging.info("Instalación de dependencias completada.")
+    subprocess.run(['nvcc', '--version'], check=True)
+
+# Llama a la función de instalación solo una vez al iniciar
 install_dependencies()
 
 import torch
@@ -71,14 +90,13 @@ MAX_SEED = np.iinfo(np.int32).max
 
 logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
 
-# Carga de modelo de Texto a Imagen (FLUX)
+# Carga de modelo de Texto a Imagen (FLUX - versión rápida)
 logging.info("Cargando pipeline de Texto a Imagen: FLUX.1-schnell")
 flux_pipeline = DiffusionPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", 
     torch_dtype=torch_dtype, 
     use_safetensors=True
-)
-flux_pipeline.to(device)
+).to(device)
 logging.info("Pipeline FLUX cargado.")
 
 # Carga de modelos de Step1X-3D
@@ -90,7 +108,7 @@ geometry_model = Step1X3DGeometryPipeline.from_pretrained(
 logging.info(f"Cargando modelo de textura: {args.texture_model}")
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
 
-logging.info("Todos los modelos han sido cargados correctamente.")
+logging.info("Todos los modelos han sido cargados correctamente y están listos.")
 
 
 # ==============================================================================
@@ -108,16 +126,15 @@ def generate_image(prompt: str, randomize_seed: bool, seed: int):
     
     generator = torch.Generator(device=device).manual_seed(int(seed))
 
-    # Prompt optimizado para obtener buenos resultados para 3D
     final_prompt = f"wbgmsst, professional 3d model {prompt}, octane render, highly detailed, volumetric, dramatic lighting, white background"
-    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting, text, watermark, signature"
+    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting, text, watermark, signature, jpeg artifacts"
 
     logging.info(f"Generando imagen con FLUX. Seed: {seed}, Prompt: '{final_prompt}'")
     
     image = flux_pipeline(
         prompt=final_prompt,
         negative_prompt=negative_prompt,
-        num_inference_steps=28, # Valor óptimo para FLUX.1-schnell
+        num_inference_steps=28,
         guidance_scale=7.5,
         generator=generator,
     ).images[0]
@@ -137,7 +154,6 @@ def generate_geometry(input_image_path, guidance_scale, inference_steps, max_fac
         
     logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
     
-    # La lógica de generación de geometría permanece igual
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -197,14 +213,12 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
     gr.Markdown("# Step1X-3D: Flujo de Texto a 3D")
     gr.Markdown("Flujo de trabajo en 3 pasos: **0. Generar Imagen → 1. Generar Geometría → 2. Generar Textura**")
     
-    # Estados para mantener las rutas de los archivos entre pasos
     image_path_state = gr.State()
     geometry_path_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=2):
-            # --- Panel de Entradas ---
-            prompt = gr.Textbox(label="Paso 0: Describe el objeto que quieres crear", value="a cute pokemon, pikachu")
+            prompt = gr.Textbox(label="Paso 0: Describe el objeto", value="a comfortable armchair")
             
             with gr.Accordion(label="Opciones Avanzadas", open=False):
                 seed = gr.Slider(0, MAX_SEED, label="Seed (para Imagen 2D)", value=42, step=1)
@@ -224,7 +238,6 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
                 btn_tex = gr.Button("2. Generate Texture", interactive=False)
 
         with gr.Column(scale=3):
-            # --- Panel de Salidas ---
             image_preview = gr.Image(label="Resultado de la Imagen Generada", type="filepath", interactive=False, height=400)
             geometry_preview = gr.Model3D(label="Vista Previa de la Geometría", height=400, clear_color=[0.0, 0.0, 0.0, 0.0])
             textured_preview = gr.Model3D(label="Vista Previa del Modelo Texturizado", height=400, clear_color=[0.0, 0.0, 0.0, 0.0])
@@ -243,12 +256,11 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
             current_seed = gr.Textbox(label="Seed Usada", interactive=False)
 
     # --- Lógica de la Interfaz ---
-
     def on_image_generated(path, used_seed):
-        """Callback que se ejecuta cuando la imagen 2D ha sido generada."""
         return {
             image_path_state: path,
             current_seed: used_seed,
+            btn_image: gr.update(interactive=True),
             btn_geo: gr.update(interactive=True, variant="primary"),
             btn_tex: gr.update(interactive=False),
             geometry_preview: gr.update(value=None),
@@ -256,36 +268,51 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
         }
 
     def on_geometry_generated(path):
-        """Callback que se ejecuta cuando la geometría se ha generado."""
         return {
             geometry_path_state: path,
+            btn_geo: gr.update(interactive=True, variant="secondary"),
             btn_tex: gr.update(interactive=True, variant="primary"),
         }
+        
+    def on_texture_generated():
+        return {
+            btn_tex: gr.update(interactive=True, variant="secondary")
+        }
 
     btn_image.click(
+        fn=lambda: gr.update(interactive=False), outputs=[btn_image]
+    ).then(
         fn=generate_image,
         inputs=[prompt, randomize_seed, seed],
         outputs=[image_preview, current_seed]
     ).then(
         fn=on_image_generated,
         inputs=[image_preview, current_seed],
-        outputs=[image_path_state, current_seed, btn_geo, btn_tex, geometry_preview, textured_preview]
+        outputs=[image_path_state, current_seed, btn_image, btn_geo, btn_tex, geometry_preview, textured_preview]
     )
     
     btn_geo.click(
+        fn=lambda: (gr.update(interactive=False), gr.update(interactive=False)),
+        outputs=[btn_geo, btn_tex]
+    ).then(
         fn=generate_geometry,
         inputs=[image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
         outputs=[geometry_preview]
     ).then(
         fn=on_geometry_generated,
         inputs=[geometry_preview],
-        outputs=[geometry_path_state, btn_tex]
+        outputs=[geometry_path_state, btn_geo, btn_tex]
     )
     
     btn_tex.click(
+       fn=lambda: gr.update(interactive=False), outputs=[btn_tex]
+    ).then(
         fn=generate_texture,
         inputs=[image_path_state, geometry_path_state],
         outputs=[textured_preview],
+    ).then(
+        fn=on_texture_generated,
+        outputs=[btn_tex]
     )
 
 demo.launch(ssr_mode=False)