Update app.py

app.py

@@ -1,46 +1,45 @@
+# ==============================================================================
+# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
+# ==============================================================================
 import os
 import shlex
 import spaces
 import subprocess
+import logging
 
-# --------------------------------------------------------------------------
-# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
-# Esta sección es crucial para que el Space de Hugging Face funcione.
-# Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias.
-# --------------------------------------------------------------------------
-def install_cuda_toolkit():
-    """Instala el toolkit de CUDA necesario para compilar las extensiones."""
-    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 = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
-    
-    print("Descargando CUDA Toolkit...")
-    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+# Configuración del logging para depuración
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')
+
+def install_dependencies():
+    """Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias."""
+    logging.info("Iniciando la instalación de dependencias...")
     
-    print("Instalando CUDA Toolkit...")
-    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+    # Instalar CUDA Toolkit
+    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"])
+    else:
+        logging.info("CUDA Toolkit ya está instalado.")
 
     os.environ["CUDA_HOME"] = "/usr/local/cuda"
-    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
-    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
-        os.environ["CUDA_HOME"],
-        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
-    )
+    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["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
-    print("Configuración de entorno de CUDA finalizada.")
-
-install_cuda_toolkit()
-
-print("Verificando instalación de PyTorch y versión de NVCC:")
-os.system("pip list | grep torch")
-os.system('nvcc -V')
-
-print("Compilando extensión de renderizador diferenciable...")
-os.system("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
+    
+    # Compilar extensiones personalizadas
+    logging.info("Compilando extensiones de renderizado...")
+    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')
 
-print("Instalando rasterizador personalizado...")
-subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
-print("Instalación y compilación completadas.")
+install_dependencies()
 
 import uuid
 import torch
@@ -48,58 +47,88 @@ import trimesh
 import argparse
 import numpy as np
 import gradio as gr
+from PIL import Image
+from diffusers import StableDiffusionXLPipeline
 from step1x3d_geometry.models.pipelines.pipeline import Step1X3DGeometryPipeline
-from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import (
-    Step1X3DTexturePipeline,
-)
+from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import Step1X3DTexturePipeline
 from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remove_degenerate_face
 
-# --------------------------------------------------------------------------
+# ==============================================================================
 # 2. CONFIGURACIÓN Y CARGA DE MODELOS
-# Aquí se definen los modelos a utilizar y se cargan en memoria.
-# --------------------------------------------------------------------------
+# ==============================================================================
+
 parser = argparse.ArgumentParser()
-parser.add_argument(
-    "--geometry_model", type=str, default="Step1X-3D-Geometry-Label-1300m"
-)
-parser.add_argument(
-    "--texture_model", type=str, default="Step1X-3D-Texture"
-)
+parser.add_argument("--geometry_model", type=str, default="Step1X-3D-Geometry-Label-1300m")
+parser.add_argument("--texture_model", type=str, default="Step1X-3D-Texture")
+parser.add_argument("--text_to_image_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0")
 parser.add_argument("--cache_dir", type=str, default="cache")
 args = parser.parse_args()
 
 os.makedirs(args.cache_dir, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16
 
-print(f"Cargando modelo de geometría: {args.geometry_model}...")
+logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
+
+# --- Carga de Modelos Step1X-3D ---
+logging.info(f"Cargando modelo de geometría: {args.geometry_model}")
 geometry_model = Step1X3DGeometryPipeline.from_pretrained(
     "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
 ).to(device)
-print("Modelo de geometría cargado.")
 
-print(f"Cargando modelo de textura: {args.texture_model}...")
+logging.info(f"Cargando modelo de textura: {args.texture_model}")
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
-print("Modelo de textura cargado.")
 
+# --- Carga de Modelo de Texto a Imagen ---
+logging.info(f"Cargando modelo de texto a imagen: {args.text_to_image_model}")
+text_to_image_pipe = StableDiffusionXLPipeline.from_pretrained(
+    args.text_to_image_model, torch_dtype=torch_dtype, variant="fp16", use_safetensors=True
+).to(device)
+logging.info("Todos los modelos han sido cargados correctamente.")
 
-# --------------------------------------------------------------------------
-# 3. FUNCIONES DE GENERACIÓN SEPARADAS
-# Se divide la lógica en dos funciones: una para la geometría y otra para la textura.
-# --------------------------------------------------------------------------
+# ==============================================================================
+# 3. FUNCIONES DE GENERACIÓN POR PASOS
+# ==============================================================================
 
-@spaces.GPU(duration=180)
-def generate_geometry(
-    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type, progress=gr.Progress(track_tqdm=True)
-):
+@spaces.GPU(duration=60)
+def generate_image_from_text(prompt, negative_prompt, guidance_scale, num_steps, seed):
     """
-    Función que genera únicamente la geometría a partir de la imagen de entrada.
+    Paso 0: Genera una imagen a partir de un prompt de texto.
     """
-    if input_image_path is None:
-        raise gr.Error("Por favor, sube una imagen para empezar.")
+    if not prompt:
+        raise gr.Error("Por favor, introduce un prompt.")
+
+    logging.info(f"Generando imagen desde el prompt: '{prompt}'")
+    
+    # Añadir modificadores para obtener un objeto aislado con fondo blanco
+    final_prompt = f"a 3d model of {prompt}, professional 3d model, octane render, highly detailed, volumetric, dramatic lighting, on a white background, studio lighting, 8k"
+    final_negative_prompt = f"text, watermark, blurry, low quality, deformed, ugly, {negative_prompt}"
+
+    generator = torch.Generator(device=device).manual_seed(int(seed))
+    
+    image = text_to_image_pipe(
+        prompt=final_prompt,
+        negative_prompt=final_negative_prompt,
+        num_inference_steps=int(num_steps),
+        guidance_scale=float(guidance_scale),
+        generator=generator,
+    ).images[0]
+    
+    save_name = str(uuid.uuid4())
+    image_path = f"{args.cache_dir}/{save_name}_generated_image.png"
+    image.save(image_path)
+    
+    logging.info(f"Imagen generada y guardada en: {image_path}")
+    return image_path
+
+@spaces.GPU(duration=180)
+def generate_geometry(input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type):
+    """Paso 1: Genera la geometría a partir de la imagen generada."""
+    if not input_image_path or not os.path.exists(input_image_path):
+        raise gr.Error("Primero debes generar una imagen desde el texto.")
         
-    print("Iniciando generación de geometría...")
+    logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
     
-    # Elige el pipeline adecuado basado en el nombre del modelo
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -118,135 +147,139 @@ def generate_geometry(
             max_facenum=int(max_facenum),
         )
 
-    # Guarda el resultado en un fichero temporal
-    save_name = str(uuid.uuid4())
-    geometry_save_path = f"{args.cache_dir}/{save_name}.glb"
+    save_name = os.path.basename(input_image_path).replace("_generated_image.png", "")
+    geometry_save_path = f"{args.cache_dir}/{save_name}_geometry.glb"
     geometry_mesh = out.mesh[0]
     geometry_mesh.export(geometry_save_path)
     
     torch.cuda.empty_cache()
-    print(f"Geometría guardada en: {geometry_save_path}")
-    
-    # Devuelve la ruta para mostrarla en el visor y para guardarla en el estado
-    return geometry_save_path, geometry_save_path
+    logging.info(f"Geometría guardada en: {geometry_save_path}")
+    return geometry_save_path
 
 @spaces.GPU(duration=120)
-def generate_texture(input_image_path, geometry_path, progress=gr.Progress(track_tqdm=True)):
-    """
-    Función que aplica la textura a una geometría ya generada.
-    """
+def generate_texture(input_image_path, geometry_path):
+    """Paso 2: Aplica la textura a la geometría generada."""
     if not geometry_path or not os.path.exists(geometry_path):
         raise gr.Error("Por favor, primero genera la geometría antes de texturizar.")
+    if not input_image_path or not os.path.exists(input_image_path):
+        raise gr.Error("Se necesita la imagen generada para el texturizado.")
         
-    print(f"Iniciando texturizado para la malla: {geometry_path}")
+    logging.info(f"Iniciando texturizado para la malla: {os.path.basename(geometry_path)}")
     geometry_mesh = trimesh.load(geometry_path)
     
-    # Post-procesamiento opcional de la malla antes de texturizar
+    # Post-procesamiento
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
     
-    # Llama al pipeline de texturizado
     textured_mesh = texture_model(input_image_path, geometry_mesh)
     
-    # Guarda el resultado final
-    save_name = os.path.basename(geometry_path).replace(".glb", "")
-    textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
+    save_name = os.path.basename(geometry_path).replace("_geometry.glb", "")
+    textured_save_path = f"{args.cache_dir}/{save_name}_textured.glb"
     textured_mesh.export(textured_save_path)
     
     torch.cuda.empty_cache()
-    print(f"Malla texturizada guardada en: {textured_save_path}")
-    
+    logging.info(f"Malla texturizada guardada en: {textured_save_path}")
     return textured_save_path
 
+# ==============================================================================
+# 4. INTERFAZ DE GRADIO
+# ==============================================================================
 
-# --------------------------------------------------------------------------
-# 4. INTERFAZ DE USUARIO CON GRADIO
-# Se define la apariencia y el comportamiento de la aplicación web.
-# --------------------------------------------------------------------------
-
-with gr.Blocks(title="Step1X-3D demo") as demo:
-    gr.Markdown("# Step1X-3D")
-    gr.Markdown("### Demostración de generación de modelos 3D a partir de una única imagen")
+with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {text-decoration: none !important;}") as demo:
+    gr.Markdown("# Step1X-3D: Flujo de Texto a Malla 3D Texturizada")
+    gr.Markdown("Flujo de trabajo en 3 pasos: **0. Texto a Imagen → 1. Generar Geometría → 2. Generar Textura**")
     
-    # Componente de estado: guarda la ruta de la geometría entre los pasos
+    # Estados para mantener las rutas de los archivos
+    generated_image_path_state = gr.State()
     geometry_path_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=2):
-            input_image = gr.Image(label="Imagen de Entrada", type="filepath")
-            
-            with gr.Accordion(label="Parámetros de Generación", open=True):
-                guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
-                inference_steps = gr.Slider(
-                    label="Pasos de Inferencia", minimum=1, maximum=100, value=50, step=1
-                )
-                max_facenum = gr.Number(label="Máx. Número de Caras", value="400000")
-                symmetry = gr.Radio(
-                    choices=["symmetry", "asymmetry"],
-                    label="Tipo de Simetría",
-                    value="symmetry",
-                    type="index",
-                )
-                edge_type = gr.Radio(
-                    choices=["sharp", "normal", "smooth"],
-                    label="Tipo de Borde",
-                    value="sharp",
-                    type="value",
-                )
-            
+            # --- Panel de Entradas ---
+            prompt = gr.Textbox(label="Paso 0: Describe el objeto que quieres crear", value="a medieval fantasy sword")
+
+            with gr.Accordion(label="Opciones Avanzadas", open=False):
+                gr.Markdown("### Opciones de Texto a Imagen (Paso 0)")
+                neg_prompt_image = gr.Textbox(label="Negative Prompt (Imagen)", value="low quality, blurry, text, watermark")
+                guidance_image = gr.Slider(0.1, 15.0, label="Guidance Scale (Imagen)", value=7.5, step=0.1)
+                steps_image = gr.Slider(10, 100, label="Steps (Imagen)", value=25, step=1)
+                seed = gr.Number(label="Seed", value=42, precision=0)
+
+                gr.Markdown("---")
+                gr.Markdown("### Opciones de Generación 3D (Paso 1)")
+                guidance_3d = gr.Number(label="Guidance Scale (3D)", value="7.5")
+                steps_3d = gr.Slider(label="Inference Steps (3D)", minimum=20, maximum=100, value=50, step=1)
+                max_facenum = gr.Number(label="Max Face Num", value="200000")
+                symmetry = gr.Radio(choices=["symmetry", "asymmetry"], label="Symmetry", value="symmetry", type="index")
+                edge_type = gr.Radio(choices=["sharp", "normal", "smooth"], label="Edge Type", value="sharp", type="value")
+
+            with gr.Row():
+                btn_gen_image = gr.Button("0. Generar Imagen", variant="secondary")
             with gr.Row():
-                btn_geo = gr.Button("1. Generar Geometría", variant="primary")
-                btn_tex = gr.Button("2. Generar Textura", visible=False, variant="primary")
+                btn_geo = gr.Button("1. Generar Geometría", interactive=False)
+                btn_tex = gr.Button("2. Generar Textura", interactive=False)
 
-        with gr.Column(scale=4):
-            textured_preview = gr.Model3D(label="Modelo con Textura", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
-            geometry_preview = gr.Model3D(label="Modelo (solo geometría)", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
+        with gr.Column(scale=3):
+            # --- Panel de Salidas ---
+            generated_image_preview = gr.Image(label="Resultado de Texto a Imagen", 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])
             
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
-                    ["examples/images/000.png"],
-                    ["examples/images/001.png"],
-                    ["examples/images/004.png"],
-                    ["examples/images/008.png"],
-                    ["examples/images/028.png"],
-                    ["examples/images/032.png"],
-                    ["examples/images/061.png"],
-                    ["examples/images/107.png"],
+                    ["a futuristic spaceship"],
+                    ["a cartoon style monster"],
+                    ["a red sports car"],
+                    ["a delicious hamburger"],
+                    ["a golden trophy"],
                 ],
-                inputs=[input_image],
-                cache_examples=False,
+                inputs=[prompt], cache_examples=False
             )
 
-    # --- Lógica de los botones y el flujo de la interfaz ---
+    # --- Lógica de la Interfaz ---
+
+    def on_image_generated(path):
+        """Se ejecuta cuando la imagen 2D ha sido generada."""
+        return {
+            generated_image_path_state: path,
+            btn_geo: gr.update(interactive=True, variant="primary"),
+            btn_tex: gr.update(interactive=False),
+            geometry_preview: gr.update(value=None),
+            textured_preview: gr.update(value=None),
+        }
+
+    def on_geometry_generated(path):
+        """Se ejecuta cuando la geometría 3D ha sido generada."""
+        return {
+            geometry_path_state: path,
+            btn_tex: gr.update(interactive=True, variant="primary"),
+        }
+
+    btn_gen_image.click(
+        fn=generate_image_from_text,
+        inputs=[prompt, neg_prompt_image, guidance_image, steps_image, seed],
+        outputs=[generated_image_preview]
+    ).then(
+        fn=on_image_generated,
+        inputs=[generated_image_preview],
+        outputs=[generated_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview]
+    )
     
-    # 1. Cuando el usuario hace clic en "Generar Geometría"
     btn_geo.click(
         fn=generate_geometry,
-        inputs=[
-            input_image,
-            guidance_scale,
-            inference_steps,
-            max_facenum,
-            symmetry,
-            edge_type,
-        ],
-        outputs=[geometry_preview, geometry_path_state]
+        inputs=[generated_image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
+        outputs=[geometry_preview]
     ).then(
-        # 2. Cuando la geometría termina, se ejecuta esta parte
-        fn=lambda: {
-            btn_tex: gr.update(visible=True),      # Hace visible el botón de texturizado
-            textured_preview: gr.update(value=None) # Limpia la vista previa de texturas anterior
-        },
-        outputs=[btn_tex, textured_preview]
+        fn=on_geometry_generated,
+        inputs=[geometry_preview],
+        outputs=[geometry_path_state, btn_tex]
     )
     
-    # 3. Cuando el usuario hace clic en "Generar Textura"
     btn_tex.click(
         fn=generate_texture,
-        inputs=[input_image, geometry_path_state],
+        inputs=[generated_image_path_state, geometry_path_state],
         outputs=[textured_preview],
     )
 
-# Lanza la aplicación
 demo.launch(ssr_mode=False)