Update app.py

app.py

@@ -1,163 +1,159 @@
-# ==============================================================================
-# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
-# ==============================================================================
+# app.py
+# Fusion: Texto -> Imagen -> Geometría -> Textura
+# Requisitos (resumido):
+# - diffusers / FluxPipeline o el pipeline que uses para generar la imagen
+# - step1x3d_geometry y step1x3d_texture (pipelines que ya usabas)
+# - trellis si lo usas (opcional)
+# Ajusta nombres de modelos, tokens y paths según tu entorno.
+
 import os
-import shlex
-import spaces
-import subprocess
-import logging
-import random
 import uuid
+import logging
+import shutil
+from typing import Tuple, Union
 
-# 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 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 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.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.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 con manejo de errores
-    logging.info("Compilando extensión 'differentiable_renderer'...")
-    renderer_path = "/home/user/app/step1x3d_texture/differentiable_renderer/"
-    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
-import trimesh
-import argparse
-import numpy as np
 import gradio as gr
 from PIL import Image
-from diffusers import DiffusionPipeline
+import numpy as np
+import torch
+import imageio
+
+# Si tienes decoradores de spaces definidos (como spaces.GPU), impórtalos.
+# from spaces import GPU  # si usas spaces.GPU
+import spaces  # si lo necesitas por compatibilidad con tus decoradores
+
+# Importa tus pipelines Step1X (o los que uses)
 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
-# ==============================================================================
+# Si usas FluxPipeline/FluxTransformer como en tu ejemplo:
+from diffusers import DiffusionPipeline  # fallback genérico; puedes dejar FluxPipeline si lo tienes
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - APP - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# -------- CONFIG (ajusta) ----------
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN", None)
+# Model names (ajusta a los subfolders que usas)
+GEOMETRY_SUBFOLDER = "Step1X-3D-Geometry-Label-1300m"
+TEXTURE_SUBFOLDER = "Step1X-3D-Texture"
+STEP1X_MODEL_REPO = "stepfun-ai/Step1X-3D"  # repo base
+# Para text-to-image: usa tu modelo preferido; aquí dejo un placeholder
+IMAGE_GEN_MODEL = "camenduru/FLUX.1-dev-diffusers"  # si usas Flux o diffusers
+# Folder para caché/temporales (por sesión)
+TMP_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
+os.makedirs(TMP_ROOT, exist_ok=True)
+# -----------------------------------
 
-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("--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
-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 - 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
-).to(device)
-logging.info("Pipeline FLUX cargado.")
-
-# Carga de modelos de 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)
-
-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 y están listos.")
-
-
-# ==============================================================================
-# 3. FUNCIONES DE GENERACIÓN POR PASOS
-# ==============================================================================
-
-@spaces.GPU(duration=60)
-def generate_image(prompt: str, randomize_seed: bool, seed: int):
-    """Paso 0: Genera una imagen a partir de un texto usando FLUX."""
-    if not prompt:
-        raise gr.Error("Por favor, introduce un prompt.")
-
+logger.info(f"Device: {device}")
+
+# ---------- Inicialización de modelos (se hace en __main__) -----------
+geometry_model = None
+texture_model = None
+image_gen_pipeline = None
+
+# ---------- FUNCIONES ----------
+def start_session(req: gr.Request):
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_ROOT, session_hash)
+    logger.info(f"[{session_hash}] start_session -> {user_dir}")
+    os.makedirs(user_dir, exist_ok=True)
+
+def end_session(req: gr.Request):
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_ROOT, session_hash)
+    logger.info(f"[{session_hash}] end_session -> {user_dir}")
+    if os.path.exists(user_dir):
+        try:
+            shutil.rmtree(user_dir)
+            logger.info(f"[{session_hash}] user dir removed")
+        except Exception as e:
+            logger.warning(f"[{session_hash}] failed removing user dir: {e}")
+
+def save_pil_image_for_session(img: Image.Image, req: gr.Request, name="generated.png") -> str:
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_ROOT, session_hash)
+    os.makedirs(user_dir, exist_ok=True)
+    path = os.path.join(user_dir, name)
+    img.save(path)
+    return path
+
+# ---------- Generar imagen desde prompt ----------
+@spaces.GPU  # si usas spaces.GPU; si no, puedes quitar
+def generate_image_from_text(
+    prompt: str,
+    seed: int,
+    randomize_seed: bool,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    req: gr.Request,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),
+) -> Image.Image:
+    """Genera una imagen 2D desde prompt y devuelve PIL.Image"""
+    global image_gen_pipeline
+    session_hash = str(req.session_hash)
+    logger.info(f"[{session_hash}] Generando imagen desde texto: '{prompt[:80]}'")
     if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
-    generator = torch.Generator(device=device).manual_seed(int(seed))
-
-    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, 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,
-        guidance_scale=7.5,
+        seed = int(np.random.randint(0, np.iinfo(np.int32).max))
+        logger.info(f"[{session_hash}] seed aleatorio -> {seed}")
+    generator = torch.Generator(device=device).manual_seed(seed) if device == "cuda" else torch.manual_seed(seed)
+
+    # Ajusta la llamada según el pipeline que uses (FluxPipeline, DiffusionPipeline, etc.)
+    # Ejemplo genérico con DiffusionPipeline (puede necesitar rename de args)
+    result = image_gen_pipeline(
+        prompt,
+        guidance_scale=float(guidance_scale),
+        num_inference_steps=8,
+        width=int(width),
+        height=int(height),
         generator=generator,
-    ).images[0]
-    
-    save_name = str(uuid.uuid4())
-    image_path = f"{args.cache_dir}/{save_name}_flux_generated.png"
-    image.save(image_path)
-    
-    logging.info(f"Imagen generada y guardada en: {image_path}")
-    return image_path, seed
+    )
+    image = result.images[0] if hasattr(result, "images") else result  # compat
+    # Guarda la imagen en carpeta de sesión
+    path = save_pil_image_for_session(image, req, name="generated_2d_image.png")
+    logger.info(f"[{session_hash}] imagen guardada en: {path}")
+    return image
 
+# ---------- Generar geometría a partir de imagen ----------
 @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 procesada."""
-    if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Primero debes generar una imagen.")
-        
-    logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
-    
-    if "Label" in args.geometry_model:
+def generate_geometry(
+    input_image: Union[str, Image.Image],
+    guidance_scale,
+    inference_steps,
+    max_facenum,
+    symmetry,
+    edge_type,
+    req: gr.Request,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),
+):
+    """
+    Genera la geometría usando Step1X geometry pipeline.
+    input_image puede ser path (str) o PIL.Image.
+    Devuelve (geometry_preview_path, geometry_path_state)
+    """
+    global geometry_model
+    session_hash = str(req.session_hash)
+    logger.info(f"[{session_hash}] Iniciando generación de geometría...")
+
+    if isinstance(input_image, str):
+        image_input = input_image
+    else:
+        # PIL.Image
+        image_input = save_pil_image_for_session(input_image, req, name="for_geometry.png")
+
+    if input_image is None:
+        raise gr.Error("Por favor, sube o genera una imagen antes de generar la geometría.")
+
+    # Lógica adaptada según si el modelo espera label u otros args
+    if "Label" in GEOMETRY_SUBFOLDER:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
-            input_image_path,
+            image_input,
             label={"symmetry": symmetry_values[int(symmetry)], "edge_type": edge_type},
             guidance_scale=float(guidance_scale),
             octree_resolution=384,
@@ -166,153 +162,161 @@ def generate_geometry(input_image_path, guidance_scale, inference_steps, max_fac
         )
     else:
         out = geometry_model(
-            input_image_path,
+            image_input,
             guidance_scale=float(guidance_scale),
             num_inference_steps=int(inference_steps),
             max_facenum=int(max_facenum),
         )
 
-    save_name = os.path.basename(input_image_path).replace("_flux_generated.png", "")
-    geometry_save_path = f"{args.cache_dir}/{save_name}_geometry.glb"
+    save_name = str(uuid.uuid4())
+    geometry_save_path = os.path.join(TMP_ROOT, session_hash, f"{save_name}.glb")
     geometry_mesh = out.mesh[0]
     geometry_mesh.export(geometry_save_path)
-    
+
     torch.cuda.empty_cache()
-    logging.info(f"Geometría guardada en: {geometry_save_path}")
-    return geometry_save_path
+    logger.info(f"[{session_hash}] Geometría guardada en: {geometry_save_path}")
+
+    # Devuelve la ruta para preview (Model3D) y para guardar en el state
+    return geometry_save_path, geometry_save_path
 
+# ---------- Generar textura a partir de geometría ----------
 @spaces.GPU(duration=120)
-def generate_texture(input_image_path, geometry_path):
-    """Paso 2: Aplica la textura a la geometría generada."""
+def generate_texture(input_image: Union[str, Image.Image], geometry_path: str, req: gr.Request, progress: gr.Progress = gr.Progress(track_tqdm=True)):
+    global texture_model
+    session_hash = str(req.session_hash)
+    logger.info(f"[{session_hash}] Iniciando texturizado para: {geometry_path}")
+
     if not geometry_path or not os.path.exists(geometry_path):
-        raise gr.Error("Por favor, primero genera la geometría.")
-    if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Se necesita la imagen generada para el texturizado.")
-        
-    logging.info(f"Iniciando texturizado para la malla: {os.path.basename(geometry_path)}")
+        raise gr.Error("Por favor, primero genera la geometría antes de texturizar.")
+
+    if isinstance(input_image, str):
+        img_path = input_image
+    else:
+        img_path = save_pil_image_for_session(input_image, req, name="for_texture.png")
+
+    # Carga y postprocesado
+    import trimesh
     geometry_mesh = trimesh.load(geometry_path)
-    
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
-    
-    textured_mesh = texture_model(input_image_path, geometry_mesh)
-    
-    save_name = os.path.basename(geometry_path).replace("_geometry.glb", "")
-    textured_save_path = f"{args.cache_dir}/{save_name}_textured.glb"
+
+    textured_mesh = texture_model(img_path, geometry_mesh)
+
+    save_name = os.path.basename(geometry_path).replace(".glb", "")
+    textured_save_path = os.path.join(TMP_ROOT, session_hash, f"{save_name}-textured.glb")
     textured_mesh.export(textured_save_path)
-    
+
     torch.cuda.empty_cache()
-    logging.info(f"Malla texturizada guardada en: {textured_save_path}")
+    logger.info(f"[{session_hash}] Malla texturizada guardada en: {textured_save_path}")
+
     return textured_save_path
 
-# ==============================================================================
-# 4. INTERFAZ DE GRADIO
-# ==============================================================================
+# ---------- Interfaz Gradio ----------
+with gr.Blocks(title="Text → Image → 3D (Step1X Flow)") as demo:
+    gr.Markdown("# Text → Image → 3D (Step1X) - Demo integrada")
+    gr.Markdown("Flujo: Texto → Generar imagen → Generar geometría → Texturizar")
 
-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 3D")
-    gr.Markdown("Flujo de trabajo en 3 pasos: **0. Generar Imagen → 1. Generar Geometría → 2. Generar Textura**")
-    
-    image_path_state = gr.State()
     geometry_path_state = gr.State()
+    generated_image_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=2):
-            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)
+            prompt = gr.Textbox(label="Prompt", placeholder="Describe el asset que quieres generar")
+            with gr.Accordion("Image Generation Settings", open=False):
+                seed = gr.Slider(0, int(2**31-1), label="Seed", value=42, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                gr.Markdown("---")
-                gr.Markdown("### Opciones de Generación 3D (Pasos 1 y 2)")
-                guidance_3d = gr.Number(label="Guidance Scale (3D)", value="7.5")
-                steps_3d = gr.Slider(label="Inference Steps (3D)", minimum=1, maximum=100, value=50)
-                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_image = gr.Button("0. Generate Image", variant="primary")
-            with gr.Row():
-                btn_geo = gr.Button("1. Generate Geometry", interactive=False)
-                btn_tex = gr.Button("2. Generate Texture", interactive=False)
+                with gr.Row():
+                    width = gr.Slider(256, 1024, label="Width", value=512, step=64)
+                    height = gr.Slider(256, 1024, label="Height", value=512, step=64)
+                guidance_scale = gr.Slider(0.0, 10.0, label="Guidance Scale", value=3.5, step=0.1)
+
+            generate_image_btn = gr.Button("1. Generar Imagen")
+            generate_geo_btn = gr.Button("2. Generar Geometría", interactive=False, visible=True)
+            generate_tex_btn = gr.Button("3. Generar Textura", interactive=False, visible=False)
 
         with gr.Column(scale=3):
-            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])
-            
+            generated_image = gr.Image(label="Imagen generada (2D)", type="pil")
+            geometry_preview = gr.Model3D(label="Geometría (GLB)", height=360)
+            textured_preview = gr.Model3D(label="Modelo texturizado (GLB)", height=360)
+
         with gr.Column(scale=1):
-            gr.Examples(
-                examples=[
-                    ["a futuristic spaceship"],
-                    ["a cartoon style monster"],
-                    ["a red sports car"],
-                    ["a medieval sword"],
-                    ["a vintage camera"],
-                ],
-                inputs=[prompt], cache_examples=False
-            )
-            current_seed = gr.Textbox(label="Seed Usada", interactive=False)
-
-    # --- Lógica de la Interfaz ---
-    def on_image_generated(path, used_seed):
-        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),
-            textured_preview: gr.update(value=None),
-        }
-
-    def on_geometry_generated(path):
-        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]
+            gr.Markdown("**Parámetros Geometría**")
+            guidance_geom = gr.Number(label="Guidance Scale geom", value=7.5)
+            inference_steps_geom = gr.Slider(1, 100, label="Pasos inferencia geom", value=50)
+            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")
+
+    # Session handlers
+    demo.load(start_session)
+    demo.unload(end_session)
+
+    # 1) Generar imagen desde texto
+    generate_image_btn.click(
+        fn=generate_image_from_text,
+        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale],
+        outputs=[generated_image],
     ).then(
-        fn=on_image_generated,
-        inputs=[image_preview, current_seed],
-        outputs=[image_path_state, current_seed, btn_image, btn_geo, btn_tex, geometry_preview, textured_preview]
+        # cuando la imagen esté lista, habilitar el botón de generar geometría
+        lambda: gr.update(interactive=True),
+        outputs=[generate_geo_btn],
     )
-    
-    btn_geo.click(
-        fn=lambda: (gr.update(interactive=False), gr.update(interactive=False)),
-        outputs=[btn_geo, btn_tex]
-    ).then(
+
+    # 2) Generar geometría desde la imagen generada
+    # usamos generated_image (PIL) y enviamos a generate_geometry
+    generate_geo_btn.click(
         fn=generate_geometry,
-        inputs=[image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
-        outputs=[geometry_preview]
+        inputs=[
+            generated_image,  # PIL image
+            guidance_geom,
+            inference_steps_geom,
+            max_facenum,
+            symmetry,
+            edge_type,
+        ],
+        outputs=[geometry_preview, geometry_path_state],
     ).then(
-        fn=on_geometry_generated,
-        inputs=[geometry_preview],
-        outputs=[geometry_path_state, btn_geo, btn_tex]
+        # habilitar el botón de generar textura
+        lambda: (gr.update(interactive=True), gr.update(visible=True)),
+        outputs=[generate_tex_btn, textured_preview]
     )
-    
-    btn_tex.click(
-       fn=lambda: gr.update(interactive=False), outputs=[btn_tex]
-    ).then(
+
+    # 3) Texturizar la geometría
+    generate_tex_btn.click(
         fn=generate_texture,
-        inputs=[image_path_state, geometry_path_state],
+        inputs=[generated_image, geometry_path_state],
         outputs=[textured_preview],
-    ).then(
-        fn=on_texture_generated,
-        outputs=[btn_tex]
     )
 
-demo.launch(ssr_mode=False)
+# ---------- Carga de modelos en main ----------
+if __name__ == "__main__":
+    # --------- Inicializar image generation pipeline ----------
+    try:
+        # Si tienes un pipeline específico (FluxPipeline) reemplaza la línea siguiente
+        logger.info("Inicializando pipeline de generación de imágenes...")
+        image_gen_pipeline = DiffusionPipeline.from_pretrained(IMAGE_GEN_MODEL, use_auth_token=HUGGINGFACE_TOKEN)
+        image_gen_pipeline = image_gen_pipeline.to(device)
+        logger.info("Pipeline de imagen cargado.")
+    except Exception as e:
+        logger.error(f"Error cargando pipeline de imágenes: {e}")
+        image_gen_pipeline = None
+
+    # --------- Inicializar Step1X modelos ----------
+    try:
+        logger.info("Cargando modelo de geometría Step1X...")
+        geometry_model = Step1X3DGeometryPipeline.from_pretrained(STEP1X_MODEL_REPO, subfolder=GEOMETRY_SUBFOLDER).to(device)
+        logger.info("Modelo de geometría cargado.")
+    except Exception as e:
+        logger.error(f"Error cargando modelo de geometría: {e}")
+        geometry_model = None
+
+    try:
+        logger.info("Cargando modelo de textura Step1X...")
+        texture_model = Step1X3DTexturePipeline.from_pretrained(STEP1X_MODEL_REPO, subfolder=TEXTURE_SUBFOLDER)
+        logger.info("Modelo de textura cargado.")
+    except Exception as e:
+        logger.error(f"Error cargando modelo de textura: {e}")
+        texture_model = None
+
+    # Lanzar app
+    demo.launch(show_error=True)