Scketch3D

app.py

@@ -2,22 +2,11 @@ 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"
@@ -27,21 +16,14 @@ def install_cuda_toolkit():
         "" 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...")
+print("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
 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 time
 import uuid
 import torch
 import trimesh
@@ -54,10 +36,7 @@ from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import (
 )
 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"
@@ -69,37 +48,19 @@ 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.")
+).to("cuda")
 
-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)
+@spaces.GPU(duration=240)
+def generate_func(
+    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
 ):
-    """
-    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
+    # geometry_model = geometry_model.to("cuda")
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -118,90 +79,49 @@ def generate_geometry(
             max_facenum=int(max_facenum),
         )
 
-    # Save the result to a temporary file
     save_name = str(uuid.uuid4())
+    print(save_name)
     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
 
+    torch.cuda.empty_cache()
+    print("Generate finish")
+    return geometry_save_path, 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")
-
+            input_image = gr.Image(label="Image", type="filepath")
+            guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
+            inference_steps = gr.Slider(
+                label="Inferece Steps", minimum=1, maximum=100, value=50
+            )
+            max_facenum = gr.Number(label="Max Face Num", 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",
+            )
+            btn = gr.Button("Start")
         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])
-            
+            textured_preview = gr.Model3D(label="Textured", height=380)
+            geometry_preview = gr.Model3D(label="Geometry", height=380)
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
@@ -218,11 +138,8 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                 cache_examples=False,
             )
 
-    # --- Button logic and interface flow ---
-    
-    # 1. When the user clicks "Generate Geometry"
-    btn_geo.click(
-        fn=generate_geometry,
+    btn.click(
+        generate_func,
         inputs=[
             input_image,
             guidance_scale,
@@ -231,22 +148,7 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
             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],
+        outputs=[geometry_preview, textured_preview],
     )
 
-# Launch the app
-demo.launch(ssr_mode=False)
+demo.launch(ssr_mode=False)