cog integration

Author: ndreca

.github/workflows/push.yml

@@ -0,0 +1,42 @@
+name: Push to Replicate
+
+on:
+  # Allows manual triggering from GitHub.com
+  workflow_dispatch:
+    inputs:
+      model_name:
+        description: "Model name to push (default: ndreca/hunyuan3d-2)"
+        required: false
+        default: "ndreca/hunyuan3d-2"
+  # Uncomment the lines below to trigger on every push to main
+  # push:
+  #   branches:
+  #     - main
+
+jobs:
+  push_to_replicate:
+    name: Push to Replicate
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Free disk space
+        uses: jlumbroso/free-disk-space@v1.3.1
+        with:
+          tool-cache: false
+          docker-images: false
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup Cog
+        uses: replicate/setup-cog@v2
+        with:
+          token: ${{ secrets.REPLICATE_API_TOKEN }}
+
+      - name: Push to Replicate
+        run: |
+          if [ -n "${{ github.event.inputs.model_name }}" ]; then
+            cog push r8.im/${{ github.event.inputs.model_name }}
+          else
+            cog push
+          fi

.gitignore

@@ -166,3 +166,9 @@ gradio_cache/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+__pycache__
+.cog
+checkpoints
+output
+.DS_Store

cog.yaml

@@ -0,0 +1,52 @@
+build:
+  gpu: true
+  cuda: "12.4"
+  python_version: "3.10"
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+    - "libglu1-mesa"
+    - "libglew2.2"
+  python_packages:
+    - gradio_litmodel3d
+    - ninja
+    - pybind11
+    - trimesh
+    - diffusers
+    - tqdm
+    - einops
+    - opencv-python
+    - numpy
+    - "torch==2.4.0"
+    - transformers
+    - torchvision
+    - torchaudio
+    - ConfigArgParse
+    - xatlas
+    - scikit-learn
+    - scikit-image
+    - tritonclient
+    - gevent
+    - geventhttpclient
+    - facexlib
+    - accelerate
+    - ipdb
+    - omegaconf
+    - pymeshlab
+    - pytorch_lightning
+    - taming-transformers-rom1504
+    - kornia
+    - rembg
+    - onnxruntime
+    - pygltflib
+    - sentencepiece
+    - gradio
+    - uvicorn
+    - "fastapi==0.112.2"
+    - wheel
+
+  run:
+    - 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

@@ -0,0 +1,121 @@
+from cog import BasePredictor, BaseModel, Input, Path
+import os
+import time
+import torch
+import subprocess
+from PIL import Image
+import shutil
+from hy3dgen.shapegen import (
+    FaceReducer, FloaterRemover, DegenerateFaceRemover,
+    Hunyuan3DDiTFlowMatchingPipeline
+)
+from hy3dgen.rembg import BackgroundRemover
+from hy3dgen.texgen import Hunyuan3DPaintPipeline
+
+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"
+
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-xf", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+class Output(BaseModel):
+    mesh: Path
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        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()
+        self.rmbg_worker = BackgroundRemover()
+
+    def predict(
+        self,
+        image: Path = Input(
+            description="Input image for generating 3D shape",
+            default=None
+        ),
+        steps: int = Input(
+            description="Number of inference steps",
+            default=50,
+            ge=20,
+            le=50,
+        ),
+        guidance_scale: float = Input(
+            description="Guidance scale for generation",
+            default=5.5,
+            ge=1.0,
+            le=20.0,
+        ),
+        seed: int = Input(
+            description="Random seed for generation",
+            default=1234
+        ),
+        octree_resolution: int = Input(
+            description="Octree resolution for mesh generation",
+            choices=[256, 384, 512],
+            default=256
+        ),
+        remove_background: bool = Input(
+            description="Whether to remove background from input image",
+            default=True
+        ),
+    ) -> Output:
+        if os.path.exists("output"):
+            shutil.rmtree("output")
+        
+        os.makedirs("output", exist_ok=True)
+
+        max_facenum = 40000
+
+        generator = torch.Generator()
+        generator = generator.manual_seed(seed)
+
+        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")
+
+        input_image.save("output/input.png")
+
+        mesh = self.i23d_worker(
+            image=input_image,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            octree_resolution=octree_resolution
+        )[0]
+
+        mesh = self.floater_remove_worker(mesh)
+        mesh = self.degenerate_face_remove_worker(mesh)
+        mesh = self.face_reduce_worker(mesh, max_facenum=max_facenum)
+        mesh = self.texgen_worker(mesh, input_image)
+        output_path = Path("output/mesh.glb")
+        mesh.export(str(output_path), include_normals=True)
+
+        if not Path(output_path).exists():
+            raise RuntimeError(f"Failed to generate mesh file at {output_path}")
+
+        return Output(mesh=output_path)