cleanup

Author: ndreca

cog.yaml

@@ -46,7 +46,7 @@ build:
     - wheel
 
   run:
-    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.9.1/pget_Linux_x86_64" && chmod +x /usr/local/bin/pget
+    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_Linux_x86_64" && chmod +x /usr/local/bin/pget
     - curl -o /tmp/custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl -L "https://huggingface.co/spaces/tencent/Hunyuan3D-2/resolve/main/custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl" && pip install /tmp/custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl
 
 predict: "predict.py:Predictor"

demo.png

@@ -9,12 +9,13 @@
     FaceReducer, FloaterRemover, DegenerateFaceRemover,
     Hunyuan3DDiTFlowMatchingPipeline
 )
+from hy3dgen.rembg import BackgroundRemover
 from hy3dgen.texgen import Hunyuan3DPaintPipeline
 
-MODEL_PATH = "checkpoints"
+MODEL_REPO = "tencent/Hunyuan3D-2"
 MODEL_URL = "https://weights.replicate.delivery/default/tencent/Hunyuan3D-2/hunyuan3d-dit-v2-0/model.tar"
-# DELIGHT_URL = "https://weights.replicate.delivery/default/tencent/Hunyuan3D-2/hunyuan3d-dit-v2-0/delight.tar"
-# PAINT_URL = "https://weights.replicate.delivery/default/tencent/Hunyuan3D-2/hunyuan3d-dit-v2-0/paint.tar"
+DELIGHT_URL = "https://weights.replicate.delivery/default/tencent/Hunyuan3D-2/hunyuan3d-dit-v2-0/delight.tar"
+PAINT_URL = "https://weights.replicate.delivery/default/tencent/Hunyuan3D-2/hunyuan3d-dit-v2-0/paint.tar"
 
 def download_weights(url, dest):
     start = time.time()
@@ -28,26 +29,24 @@ class Output(BaseModel):
 
 class Predictor(BasePredictor):
     def setup(self) -> None:
-        """Load the model into memory to make running multiple predictions efficient"""
-        #make directory model_path
-        os.makedirs(MODEL_PATH, exist_ok=True)
-
-        # Download weights
-        if not os.path.exists(MODEL_PATH + "/hunyuan3d-dit-v2-0"):
-            download_weights(MODEL_URL, MODEL_PATH + "/hunyuan3d-dit-v2-0")
-        # if not os.path.exists(MODEL_PATH + "/hunyuan3d-delight-v2-0"):
-        #     download_weights(DELIGHT_URL, MODEL_PATH + "/hunyuan3d-delight-v2-0")
-        # if not os.path.exists(MODEL_PATH + "/hunyuan3d-paint-v2-0"):
-        #     download_weights(PAINT_URL, MODEL_PATH + "/hunyuan3d-paint-v2-0")
-
-        # Load shape generation model
-        self.pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(MODEL_PATH+"/hunyuan3d-dit-v2-0")
+        model_base_path = os.path.expanduser(os.path.join("~/.cache/hy3dgen", MODEL_REPO))
+        os.makedirs(model_base_path, exist_ok=True)
+
+        model_path = os.path.join(model_base_path, "hunyuan3d-dit-v2-0")
+        delight_path = os.path.join(model_base_path, "hunyuan3d-delight-v2-0")
+        paint_path = os.path.join(model_base_path, "hunyuan3d-paint-v2-0")
+        if not os.path.exists(model_path):
+            download_weights(MODEL_URL, model_path)
+        if not os.path.exists(delight_path):
+            download_weights(DELIGHT_URL, delight_path)
+        if not os.path.exists(paint_path):
+            download_weights(PAINT_URL, paint_path)
+
+        self.i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(MODEL_REPO)
+        self.texgen_worker = Hunyuan3DPaintPipeline.from_pretrained(MODEL_REPO)
         self.floater_remove_worker = FloaterRemover()
         self.degenerate_face_remove_worker = DegenerateFaceRemover()
         self.face_reduce_worker = FaceReducer()
-
-        # Load background remover
-        from hy3dgen.rembg import BackgroundRemover
         self.rmbg_worker = BackgroundRemover()
 
     def predict(
@@ -82,51 +81,39 @@ def predict(
             default=True
         ),
     ) -> Output:
-        """Run a single prediction on the model"""
-        # Clean up past runs, delete output directory
         if os.path.exists("output"):
             shutil.rmtree("output")
 
-        # Create output directory
         os.makedirs("output", exist_ok=True)
-        
-        # Set random seed
+
         generator = torch.Generator()
         generator = generator.manual_seed(seed)
 
-        # Process input image
         if image is not None:
             input_image = Image.open(str(image))
             if remove_background or input_image.mode == "RGB":
                 input_image = self.rmbg_worker(input_image.convert('RGB'))
         else:
             raise ValueError("Image must be provided")
 
-        # Save processed input image
         input_image.save("output/input.png")
 
-        # Generate 3D mesh
-        mesh = self.pipeline(
+        mesh = self.i23d_worker(
             image=input_image,
             num_inference_steps=steps,
             guidance_scale=guidance_scale,
             generator=generator,
             octree_resolution=octree_resolution
         )[0]
 
-        # Post-process mesh
         mesh = self.floater_remove_worker(mesh)
         mesh = self.degenerate_face_remove_worker(mesh)
         mesh = self.face_reduce_worker(mesh)
+        textured_mesh = self.texgen_worker(mesh, input_image)
 
-        # Set mesh color to dark gray
-        mesh.visual.face_colors = [64, 64, 64, 255]  # RGB + Alpha for dark gray
-
-        # Export mesh
-        output_path = Path("output/gray_mesh.glb")
-        mesh.export(str(output_path), include_normals=False)
+        output_path = Path("output/mesh.glb")
+        textured_mesh.export(str(output_path), include_normals=True)
 
-        # Verify file exists and is readable
         if not Path(output_path).exists():
             raise RuntimeError(f"Failed to generate mesh file at {output_path}")