# ============================================================================== # 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS # ============================================================================== import os import shlex import spaces import subprocess import logging # Configuración del logging para depuración 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.""" logging.info("Iniciando la instalación de dependencias...") 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"] = 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" 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') install_dependencies() import uuid import torch import trimesh import argparse import random import numpy as np import gradio as gr from PIL import Image from huggingface_hub import hf_hub_download 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 # ============================================================================== # 2. CONFIGURACIÓN Y CARGA DE MODELOS # ============================================================================== 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.bfloat16 MAX_SEED = np.iinfo(np.int32).max 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) 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 FLUX para Texto-a-Imagen --- logging.info("Cargando modelo FLUX para Texto-a-Imagen...") single_file_base_model = "camenduru/FLUX.1-dev-diffusers" # --- CORRECCIÓN AQUÍ --- # Descargar el archivo GGUF explícitamente usando hf_hub_download flux_repo_id = "gokaygokay/flux-game" flux_filename = "hyperflux_00001_.q8_0.gguf" logging.info(f"Descargando {flux_filename} desde {flux_repo_id}...") downloaded_flux_path = hf_hub_download(repo_id=flux_repo_id, filename=flux_filename) logging.info(f"Archivo GGUF descargado en: {downloaded_flux_path}") # --- FIN DE LA CORRECCIÓN --- 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) # Usar la ruta local descargada transformer = FluxTransformer2DModel.from_single_file(downloaded_flux_path, 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 (Sin cambios) # ============================================================================== @spaces.GPU(duration=60) 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 de texto no puede estar vacío.") if randomize_seed: seed = random.randint(0, MAX_SEED) logging.info(f"Generando imagen con prompt: '{prompt}' y seed: {seed}") generator = torch.Generator(device=device).manual_seed(seed) 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, 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}_t2i.png" image.save(image_save_path) logging.info(f"Imagen 2D generada y guardada en: {image_save_path}") 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 desde el texto.") logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}") if "Label" in args.geometry_model: symmetry_values = ["x", "asymmetry"] out = geometry_model( input_image_path, label={"symmetry": symmetry_values[int(symmetry)], "edge_type": edge_type}, guidance_scale=float(guidance_scale), octree_resolution=384, max_facenum=int(max_facenum), num_inference_steps=int(inference_steps), ) else: out = geometry_model( input_image_path, 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("_t2i.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() 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): """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.") 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) 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) torch.cuda.empty_cache() logging.info(f"Malla texturizada guardada en: {textured_save_path}") return textured_save_path # ============================================================================== # 4. INTERFAZ DE GRADIO (Sin cambios) # ============================================================================== 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**") generated_image_path_state = gr.State() geometry_path_state = gr.State() with gr.Row(): with gr.Column(scale=2): prompt = gr.Textbox(label="Paso 0: Describe tu objeto", placeholder="Ej: a treasure chest, a sci-fi helmet, a cute dog") 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) 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_t2i = gr.Button("0. Generate Image", variant="secondary") with gr.Row(): btn_geo = gr.Button("1. Generate Geometry", interactive=False) btn_tex = gr.Button("2. Generate Texture", interactive=False) with gr.Column(scale=3): 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]) def on_image_generated(path): 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): return { geometry_path_state: path, btn_tex: gr.update(interactive=True, variant="primary"), } btn_t2i.click( fn=generate_image_from_text, inputs=[prompt, seed, randomize_seed, guidance_t2i, steps_t2i], 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] ) btn_geo.click( fn=generate_geometry, inputs=[generated_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] ) btn_tex.click( fn=generate_texture, inputs=[generated_image_path_state, geometry_path_state], outputs=[textured_preview], ) demo.launch(ssr_mode=False)