Update app.py

app.py

@@ -8,7 +8,7 @@ import subprocess
 import logging
 
 # Configuración del logging para depuración
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D-T2I - %(levelname)s - %(message)s')
 
 def install_dependencies():
     """Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias."""
@@ -47,7 +47,8 @@ import random
 import numpy as np
 import gradio as gr
 from PIL import Image
-from diffusers import FluxPipeline
+from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig
+from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfigTF
 from step1x3d_geometry.models.pipelines.pipeline import Step1X3DGeometryPipeline
 from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import Step1X3DTexturePipeline
 from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remove_degenerate_face
@@ -69,16 +70,6 @@ MAX_SEED = np.iinfo(np.int32).max
 
 logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
 
-# --- Carga del modelo FLUX de Texto a Imagen desde la versión de la comunidad ---
-logging.info("Cargando modelo FLUX.1-dev desde camenduru/FLUX.1-dev-diffusers...")
-flux_pipe = FluxPipeline.from_pretrained(
-    "camenduru/FLUX.1-dev-diffusers", 
-    torch_dtype=torch_dtype,
-    variant="fp16"
-)
-flux_pipe.to(device)
-logging.info("Modelo FLUX cargado.")
-
 # --- Carga de Modelos Step1X-3D ---
 logging.info(f"Cargando modelo de geometría: {args.geometry_model}")
 geometry_model = Step1X3DGeometryPipeline.from_pretrained(
@@ -88,54 +79,64 @@ 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.")
+# --- Carga de Modelo FLUX para Texto-a-Imagen ---
+logging.info("Cargando modelo FLUX para Texto-a-Imagen...")
+single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
+file_url = "https://huggingface.co/gokaygokay/flux-game/resolve/main/hyperflux_00001_.q8_0.gguf"
 
+quantization_config_tf = BitsAndBytesConfigTF(load_in_8bit=True, bnb_8bit_compute_dtype=torch_dtype)
+text_encoder_2 = T5EncoderModel.from_pretrained(single_file_base_model, subfolder="text_encoder_2", torch_dtype=torch_dtype, quantization_config=quantization_config_tf)
+transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", quantization_config=GGUFQuantizationConfig(compute_dtype=torch_dtype), torch_dtype=torch_dtype, config=single_file_base_model)
+flux_pipeline = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch_dtype)
+flux_pipeline.to(device)
+logging.info("Todos los modelos han sido cargados correctamente.")
 
 # ==============================================================================
 # 3. FUNCIONES DE GENERACIÓN POR PASOS
 # ==============================================================================
 
 @spaces.GPU(duration=60)
-def generate_image_from_text(prompt, negative_prompt, guidance_scale, num_steps, seed, randomize_seed):
-    """
-    Paso 0: Genera una imagen 2D a partir de un texto usando FLUX.
-    """
+def generate_image_from_text(prompt, seed, randomize_seed, guidance_scale, num_steps):
+    """Paso 0: Genera una imagen 2D a partir de un prompt de texto usando FLUX."""
     if not prompt:
-        raise gr.Error("El prompt no puede estar vacío.")
-        
+        raise gr.Error("El prompt de texto no puede estar vacío.")
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    generator = torch.Generator(device=device).manual_seed(int(seed))
-
-    logging.info(f"Generando imagen con prompt: '{prompt}', Seed: {seed}")
-    
-    # Añadir modificadores para mejorar la calidad y el estilo 3D
-    final_prompt = f"3d model, {prompt}, octane render, professionally rendered, high quality, white background"
+    logging.info(f"Generando imagen con prompt: '{prompt}' y seed: {seed}")
+    generator = torch.Generator(device=device).manual_seed(seed)
     
-    image = flux_pipe(
-        prompt=final_prompt,
+    # Prompting específico que funciona bien con el modelo FLUX para objetos
+    full_prompt = f"professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting, on a white background"
+    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting, photo, text, watermark"
+
+    image = flux_pipeline(
+        prompt=full_prompt,
         negative_prompt=negative_prompt,
-        num_inference_steps=int(num_steps),
         guidance_scale=float(guidance_scale),
+        num_inference_steps=int(num_steps),
+        width=1024,
+        height=1024,
         generator=generator,
     ).images[0]
     
     save_name = str(uuid.uuid4())
-    image_save_path = f"{args.cache_dir}/{save_name}_generated_image.png"
+    image_save_path = f"{args.cache_dir}/{save_name}_t2i.png"
     image.save(image_save_path)
     
     logging.info(f"Imagen 2D generada y guardada en: {image_save_path}")
-    return image_save_path, seed
+    return image_save_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 a partir del texto.")
+        raise gr.Error("Primero debes generar una imagen desde el texto.")
         
     logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
     
+    # ... (El resto de esta función es idéntico al de la respuesta anterior)
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -154,7 +155,7 @@ def generate_geometry(input_image_path, guidance_scale, inference_steps, max_fac
             max_facenum=int(max_facenum),
         )
 
-    save_name = os.path.basename(input_image_path).replace("_generated_image.png", "")
+    save_name = os.path.basename(input_image_path).replace("_t2i.png", "")
     geometry_save_path = f"{args.cache_dir}/{save_name}_geometry.glb"
     geometry_mesh = out.mesh[0]
     geometry_mesh.export(geometry_save_path)
@@ -172,14 +173,12 @@ def generate_texture(input_image_path, geometry_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)}")
-    geometry_mesh = trimesh.load(geometry_path)
     
-    # Post-procesamiento
+    # ... (El resto de esta función es idéntico al de la respuesta anterior)
+    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.export(textured_save_path)
@@ -192,85 +191,72 @@ def generate_texture(input_image_path, geometry_path):
 # 4. INTERFAZ DE GRADIO
 # ==============================================================================
 
-with gr.Blocks(title="Step1X-3D: Texto a 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. Texto a Imagen → 1. Generar Geometría → 2. Generar Textura**")
+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. Generar Imagen → 1. Generar Geometría → 2. Generar Textura**")
     
-    # Estados para mantener las rutas de los archivos
+    # Estados para mantener las rutas de los archivos entre pasos
     generated_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 3D que quieres crear", value="a detailed, realistic battle axe")
+            prompt = gr.Textbox(label="Paso 0: Describe tu objeto", placeholder="Ej: a treasure chest, a sci-fi helmet, a cute dog")
             
-            with gr.Accordion(label="Opciones Avanzadas", open=False):
-                gr.Markdown("### Opciones de Texto a Imagen (Paso 0)")
-                neg_prompt = gr.Textbox(label="Negative Prompt (Imagen)", value="blurry, low quality, bad, text, watermark")
-                guidance_image = gr.Slider(0.0, 10.0, label="Guidance Scale (Imagen)", value=4.0, step=0.1)
-                steps_image = gr.Slider(10, 50, label="Steps (Imagen)", value=28, step=1)
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+            with gr.Accordion("Opciones de Generación de Imagen (Paso 0)", open=True):
+                seed = gr.Slider(0, MAX_SEED, label="Seed", value=42, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                guidance_t2i = gr.Slider(0.0, 10.0, label="Guidance Scale (Imagen)", value=3.5, step=0.1)
+                steps_t2i = gr.Slider(1, 20, label="Steps (Imagen)", value=8, step=1)
 
-                gr.Markdown("---")
-                gr.Markdown("### Opciones de Generación 3D (Paso 1)")
+            with gr.Accordion("Opciones de Generación 3D (Pasos 1 y 2)", open=False):
                 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_gen_image = gr.Button("0. Generar Imagen", variant="secondary")
+                btn_t2i = gr.Button("0. Generate Image", variant="secondary")
             with gr.Row():
-                btn_geo = gr.Button("1. Generar Geometría", interactive=False)
-                btn_tex = gr.Button("2. Generar Textura", interactive=False)
-        
+                btn_geo = gr.Button("1. Generate Geometry", interactive=False)
+                btn_tex = gr.Button("2. Generate Texture", interactive=False)
+
         with gr.Column(scale=3):
             # --- Panel de Salidas ---
-            image_preview = gr.Image(label="Resultado de Texto a Imagen", type="filepath", interactive=False, height=400)
+            generated_image_preview = gr.Image(label="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])
-
-        with gr.Column(scale=1):
-            gr.Examples(
-                examples=[
-                    ["a small wooden chest with gold trim"],
-                    ["a futuristic sci-fi pistol"],
-                    ["a cute, chibi-style red dragon"],
-                    ["a slice of pizza with pepperoni and mushrooms"],
-                    ["a classic leather-bound book with a gold clasp"],
-                ],
-                inputs=[prompt], cache_examples=False
-            )
-
+            
     # --- Lógica de la Interfaz ---
-    
-    def on_image_generated(path, current_seed):
+
+    def on_image_generated(path):
+        """Callback para actualizar la UI después de generar la imagen."""
         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),
-            seed: gr.update(value=int(current_seed)) # Actualiza el slider de la seed
         }
 
     def on_geometry_generated(path):
+        """Callback para actualizar la UI después de generar la geometría."""
         return {
             geometry_path_state: path,
             btn_tex: gr.update(interactive=True, variant="primary"),
         }
 
-    btn_gen_image.click(
+    # Cadena de eventos
+    btn_t2i.click(
         fn=generate_image_from_text,
-        inputs=[prompt, neg_prompt, guidance_image, steps_image, seed, randomize_seed],
-        outputs=[image_preview, seed]
+        inputs=[prompt, seed, randomize_seed, guidance_t2i, steps_t2i],
+        outputs=[generated_image_preview]
     ).then(
         fn=on_image_generated,
-        inputs=[image_preview, seed],
-        outputs=[generated_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview, seed]
+        inputs=[generated_image_preview],
+        outputs=[generated_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview]
     )
     
     btn_geo.click(