1

app.py

@@ -39,49 +39,6 @@ from step1x3d_texture.utils.shape_post_process import (
     DegenerateFaceRemover,
 )
 
-geometry_model = None
-texture_model = None
-
-@spaces.GPU(duration=240)
-def generate_func(
-    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
-):
-    geometry_model = geometry_model.to("cuda")
-    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 = 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)
-
-    geometry_mesh = DegenerateFaceRemover()(geometry_mesh)
-    geometry_mesh = FaceReducer()(geometry_mesh)
-    textured_mesh = texture_model(input_image_path, geometry_mesh)
-    textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
-    textured_mesh.export(textured_save_path)
-
-    torch.cuda.empty_cache()
-    print("Generate finish")
-    return geometry_save_path, textured_save_path
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -101,6 +58,47 @@ if __name__ == "__main__":
 
     texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
 
+
+    @spaces.GPU(duration=240)
+    def generate_func(
+        input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
+    ):
+        geometry_model = geometry_model.to("cuda")
+        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 = 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)
+
+        geometry_mesh = DegenerateFaceRemover()(geometry_mesh)
+        geometry_mesh = FaceReducer()(geometry_mesh)
+        textured_mesh = texture_model(input_image_path, geometry_mesh)
+        textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
+        textured_mesh.export(textured_save_path)
+
+        torch.cuda.empty_cache()
+        print("Generate finish")
+        return geometry_save_path, textured_save_path
+
+
     with gr.Blocks(title="Step1X-3D demo") as demo:
         gr.Markdown("# Step1X-3D")
         with gr.Row():