Update app.py

app.py

@@ -6,6 +6,8 @@ import shlex
 import spaces
 import subprocess
 import logging
+import random
+import uuid
 
 # Configuración del logging para depuración
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')
@@ -41,14 +43,13 @@ def install_dependencies():
 
 install_dependencies()
 
-import uuid
 import torch
 import trimesh
 import argparse
 import numpy as np
 import gradio as gr
 from PIL import Image
-from diffusers import StableDiffusionXLPipeline
+from diffusers import DiffusionPipeline
 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
@@ -60,17 +61,27 @@ from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remov
 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("--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
+MAX_SEED = np.iinfo(np.int32).max
 
 logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
 
-# --- Carga de Modelos Step1X-3D ---
+# Carga de modelo de Texto a Imagen (FLUX)
+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)
+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
@@ -79,56 +90,54 @@ 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)
 
-# --- 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 POR PASOS
 # ==============================================================================
 
 @spaces.GPU(duration=60)
-def generate_image_from_text(prompt, negative_prompt, guidance_scale, num_steps, seed):
-    """
-    Paso 0: Genera una imagen a partir de un prompt de texto.
-    """
+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.")
 
-    logging.info(f"Generando imagen desde el prompt: '{prompt}'")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
     
-    # 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))
+
+    # 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"
+
+    logging.info(f"Generando imagen con FLUX. Seed: {seed}, Prompt: '{final_prompt}'")
     
-    image = text_to_image_pipe(
+    image = flux_pipeline(
         prompt=final_prompt,
-        negative_prompt=final_negative_prompt,
-        num_inference_steps=int(num_steps),
-        guidance_scale=float(guidance_scale),
+        negative_prompt=negative_prompt,
+        num_inference_steps=28, # Valor óptimo para FLUX.1-schnell
+        guidance_scale=7.5,
         generator=generator,
     ).images[0]
     
     save_name = str(uuid.uuid4())
-    image_path = f"{args.cache_dir}/{save_name}_generated_image.png"
+    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
+    return image_path, seed
 
 @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."""
+    """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 desde el texto.")
+        raise gr.Error("Primero debes generar una imagen.")
         
     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(
@@ -147,7 +156,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("_flux_generated.png", "")
     geometry_save_path = f"{args.cache_dir}/{save_name}_geometry.glb"
     geometry_mesh = out.mesh[0]
     geometry_mesh.export(geometry_save_path)
@@ -160,14 +169,13 @@ def generate_geometry(input_image_path, guidance_scale, inference_steps, max_fac
 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.")
+        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)}")
     geometry_mesh = trimesh.load(geometry_path)
     
-    # Post-procesamiento
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
     
@@ -186,42 +194,38 @@ def generate_texture(input_image_path, geometry_path):
 # ==============================================================================
 
 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**")
+    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
-    generated_image_path_state = gr.State()
+    # 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 medieval fantasy sword")
-
+            prompt = gr.Textbox(label="Paso 0: Describe el objeto que quieres crear", value="a cute pokemon, pikachu")
+            
             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)
-
+                seed = gr.Slider(0, MAX_SEED, label="Seed (para Imagen 2D)", value=42, step=1)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                 gr.Markdown("---")
-                gr.Markdown("### Opciones de Generación 3D (Paso 1)")
+                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=20, maximum=100, value=50, step=1)
+                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_image = gr.Button("0. Generate Image", variant="primary")
             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 ---
-            generated_image_preview = gr.Image(label="Resultado de Texto a Imagen", type="filepath", interactive=False, height=400)
+            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])
             
@@ -231,18 +235,20 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
                     ["a futuristic spaceship"],
                     ["a cartoon style monster"],
                     ["a red sports car"],
-                    ["a delicious hamburger"],
-                    ["a golden trophy"],
+                    ["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):
-        """Se ejecuta cuando la imagen 2D ha sido generada."""
+    def on_image_generated(path, used_seed):
+        """Callback que se ejecuta cuando la imagen 2D ha sido generada."""
         return {
-            generated_image_path_state: path,
+            image_path_state: path,
+            current_seed: used_seed,
             btn_geo: gr.update(interactive=True, variant="primary"),
             btn_tex: gr.update(interactive=False),
             geometry_preview: gr.update(value=None),
@@ -250,25 +256,25 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
         }
 
     def on_geometry_generated(path):
-        """Se ejecuta cuando la geometría 3D ha sido generada."""
+        """Callback que se ejecuta cuando la geometría se ha generado."""
         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]
+    btn_image.click(
+        fn=generate_image,
+        inputs=[prompt, randomize_seed, seed],
+        outputs=[image_preview, current_seed]
     ).then(
         fn=on_image_generated,
-        inputs=[generated_image_preview],
-        outputs=[generated_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview]
+        inputs=[image_preview, current_seed],
+        outputs=[image_path_state, current_seed, btn_geo, btn_tex, geometry_preview, textured_preview]
     )
     
     btn_geo.click(
         fn=generate_geometry,
-        inputs=[generated_image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
+        inputs=[image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
         outputs=[geometry_preview]
     ).then(
         fn=on_geometry_generated,
@@ -278,7 +284,7 @@ with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {tex
     
     btn_tex.click(
         fn=generate_texture,
-        inputs=[generated_image_path_state, geometry_path_state],
+        inputs=[image_path_state, geometry_path_state],
         outputs=[textured_preview],
     )