app.py

import os
import shlex
import spaces
import subprocess

# --------------------------------------------------------------------------
# 1. ENVIRONMENT AND DEPENDENCY INSTALLATION
# This section is crucial for the Hugging Face Space to work.
# Installs the CUDA toolkit and compiles the necessary C++/CUDA extensions.
# --------------------------------------------------------------------------
def install_cuda_toolkit():
    """Installs the CUDA toolkit required to compile extensions."""
    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("Downloading CUDA Toolkit...")
    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
    
    print("Installing CUDA Toolkit...")
    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])

    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["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
    print("CUDA environment configuration completed.")

install_cuda_toolkit()

print("Verifying PyTorch installation and NVCC version:")
os.system("pip list | grep torch")
os.system('nvcc -V')

print("Compiling differentiable renderer extension...")
os.system("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")

print("Installing custom rasterizer...")
subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
print("Installation and compilation completed.")

import uuid
import torch
import trimesh
import argparse
import numpy as np
import gradio as gr
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. MODEL CONFIGURATION AND LOADING
# Here we define the models to be used and load them into memory.
# --------------------------------------------------------------------------
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"

print(f"Loading geometry model: {args.geometry_model}...")
geometry_model = Step1X3DGeometryPipeline.from_pretrained(
    "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
).to(device)
print("Geometry model loaded.")

print(f"Loading texture model: {args.texture_model}...")
texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
print("Texture model loaded.")


# --------------------------------------------------------------------------
# 3. SEPARATE GENERATION FUNCTIONS
# The logic is split into two functions: one for geometry and one for textures.
# --------------------------------------------------------------------------

@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)
):
    """
    Function that generates geometry only from the input image.
    """
    if input_image_path is None:
        raise gr.Error("Please upload an image to start.")
        
    print("Starting geometry generation...")
    
    # Choose the appropriate pipeline based on the model name
    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 the result to a temporary file
    save_name = str(uuid.uuid4())
    geometry_save_path = f"{args.cache_dir}/{save_name}.glb"
    geometry_mesh = out.mesh[0]
    geometry_mesh.export(geometry_save_path)
    
    torch.cuda.empty_cache()
    print(f"Geometry saved at: {geometry_save_path}")
    
    # Return the path for display in the viewer and to store in state
    return geometry_save_path, geometry_save_path

@spaces.GPU(duration=120)
def generate_texture(input_image_path, geometry_path, progress=gr.Progress(track_tqdm=True)):
    """
    Function that applies texture to an already generated geometry.
    """
    if not geometry_path or not os.path.exists(geometry_path):
        raise gr.Error("Please generate the geometry first before texturing.")
        
    print(f"Starting texturing for mesh: {geometry_path}")
    geometry_mesh = trimesh.load(geometry_path)
    
    # Optional post-processing of the mesh before texturing
    geometry_mesh = remove_degenerate_face(geometry_mesh)
    geometry_mesh = reduce_face(geometry_mesh)
    
    # Call the texturing pipeline
    textured_mesh = texture_model(input_image_path, geometry_mesh)
    
    # Save the final result
    save_name = os.path.basename(geometry_path).replace(".glb", "")
    textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
    textured_mesh.export(textured_save_path)
    
    torch.cuda.empty_cache()
    print(f"Textured mesh saved at: {textured_save_path}")
    
    return textured_save_path


# --------------------------------------------------------------------------
# 4. GRADIO USER INTERFACE
# Defines the look and behavior of the web app.
# --------------------------------------------------------------------------

with gr.Blocks(title="Step1X-3D demo") as demo:
    gr.Markdown("# Step1X-3D")
    gr.Markdown("### Demo for generating 3D models from a single image")
    
    # State component: stores the geometry path between steps
    geometry_path_state = gr.State()

    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Input Image", type="filepath")
            
            with gr.Accordion(label="Generation Parameters", open=True):
                guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
                inference_steps = gr.Slider(
                    label="Inference Steps", minimum=1, maximum=100, value=50, step=1
                )
                max_facenum = gr.Number(label="Max. Number of Faces", value="400000")
                symmetry = gr.Radio(
                    choices=["symmetry", "asymmetry"],
                    label="Symmetry Type",
                    value="symmetry",
                    type="index",
                )
                edge_type = gr.Radio(
                    choices=["sharp", "normal", "smooth"],
                    label="Edge Type",
                    value="sharp",
                    type="value",
                )
            
            with gr.Row():
                btn_geo = gr.Button("1. Generate Geometry", variant="primary")
                btn_tex = gr.Button("2. Generate Texture", visible=False, variant="primary")

        with gr.Column(scale=4):
            textured_preview = gr.Model3D(label="Textured Model", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
            geometry_preview = gr.Model3D(label="Model (geometry only)", height=380, 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"],
                ],
                inputs=[input_image],
                cache_examples=False,
            )

    # --- Button logic and interface flow ---
    
    # 1. When the user clicks "Generate Geometry"
    btn_geo.click(
        fn=generate_geometry,
        inputs=[
            input_image,
            guidance_scale,
            inference_steps,
            max_facenum,
            symmetry,
            edge_type,
        ],
        outputs=[geometry_preview, geometry_path_state]
    ).then(
        # 2. When geometry is done, run this part
        fn=lambda: {
            btn_tex: gr.update(visible=True),      # Show texture button
            textured_preview: gr.update(value=None) # Clear previous texture preview
        },
        outputs=[btn_tex, textured_preview]
    )
    
    # 3. When the user clicks "Generate Texture"
    btn_tex.click(
        fn=generate_texture,
        inputs=[input_image, geometry_path_state],
        outputs=[textured_preview],
    )

# Launch the app
demo.launch(ssr_mode=False)