| import base64 |
| import logging |
| import os |
| import sys |
| from contextlib import nullcontext |
|
|
| import comfy.model_management |
| import folder_paths |
| import numpy as np |
| import torch |
| import trimesh |
| from PIL import Image |
| from trimesh.exchange import gltf |
|
|
| sys.path.append(os.path.dirname(__file__)) |
| from sf3d.system import SF3D |
| from sf3d.utils import resize_foreground |
|
|
| SF3D_CATEGORY = "StableFast3D" |
| SF3D_MODEL_NAME = "stabilityai/stable-fast-3d" |
|
|
|
|
| class StableFast3DLoader: |
| CATEGORY = SF3D_CATEGORY |
| FUNCTION = "load" |
| RETURN_NAMES = ("sf3d_model",) |
| RETURN_TYPES = ("SF3D_MODEL",) |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": {}} |
|
|
| def load(self): |
| device = comfy.model_management.get_torch_device() |
| model = SF3D.from_pretrained( |
| SF3D_MODEL_NAME, |
| config_name="config.yaml", |
| weight_name="model.safetensors", |
| ) |
| model.to(device) |
| model.eval() |
|
|
| return (model,) |
|
|
|
|
| class StableFast3DPreview: |
| CATEGORY = SF3D_CATEGORY |
| FUNCTION = "preview" |
| OUTPUT_NODE = True |
| RETURN_TYPES = () |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return {"required": {"mesh": ("MESH",)}} |
|
|
| def preview(self, mesh): |
| glbs = [] |
| for m in mesh: |
| scene = trimesh.Scene(m) |
| glb_data = gltf.export_glb(scene, include_normals=True) |
| glb_base64 = base64.b64encode(glb_data).decode("utf-8") |
| glbs.append(glb_base64) |
| return {"ui": {"glbs": glbs}} |
|
|
|
|
| class StableFast3DSampler: |
| CATEGORY = SF3D_CATEGORY |
| FUNCTION = "predict" |
| RETURN_NAMES = ("mesh",) |
| RETURN_TYPES = ("MESH",) |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "model": ("SF3D_MODEL",), |
| "image": ("IMAGE",), |
| "foreground_ratio": ( |
| "FLOAT", |
| {"default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01}, |
| ), |
| "texture_resolution": ( |
| "INT", |
| {"default": 1024, "min": 512, "max": 2048, "step": 256}, |
| ), |
| }, |
| "optional": { |
| "mask": ("MASK",), |
| "remesh": (["none", "triangle", "quad"],), |
| "vertex_count": ( |
| "INT", |
| {"default": -1, "min": -1, "max": 20000, "step": 1}, |
| ), |
| }, |
| } |
|
|
| def predict( |
| s, |
| model, |
| image, |
| mask, |
| foreground_ratio, |
| texture_resolution, |
| remesh="none", |
| vertex_count=-1, |
| ): |
| if image.shape[0] != 1: |
| raise ValueError("Only one image can be processed at a time") |
|
|
| pil_image = Image.fromarray( |
| torch.clamp(torch.round(255.0 * image[0]), 0, 255) |
| .type(torch.uint8) |
| .cpu() |
| .numpy() |
| ) |
|
|
| if mask is not None: |
| print("Using Mask") |
| mask_np = np.clip(255.0 * mask[0].detach().cpu().numpy(), 0, 255).astype( |
| np.uint8 |
| ) |
| mask_pil = Image.fromarray(mask_np, mode="L") |
| pil_image.putalpha(mask_pil) |
| else: |
| if image.shape[3] != 4: |
| print("No mask or alpha channel detected, Converting to RGBA") |
| pil_image = pil_image.convert("RGBA") |
|
|
| pil_image = resize_foreground(pil_image, foreground_ratio) |
| print(remesh) |
| with torch.no_grad(): |
| with torch.autocast( |
| device_type="cuda", dtype=torch.bfloat16 |
| ) if "cuda" in comfy.model_management.get_torch_device().type else nullcontext(): |
| mesh, glob_dict = model.run_image( |
| pil_image, |
| bake_resolution=texture_resolution, |
| remesh=remesh, |
| vertex_count=vertex_count, |
| ) |
|
|
| if mesh.vertices.shape[0] == 0: |
| raise ValueError("No subject detected in the image") |
|
|
| return ([mesh],) |
|
|
|
|
| class StableFast3DSave: |
| CATEGORY = SF3D_CATEGORY |
| FUNCTION = "save" |
| OUTPUT_NODE = True |
| RETURN_TYPES = () |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "mesh": ("MESH",), |
| "filename_prefix": ("STRING", {"default": "SF3D"}), |
| } |
| } |
|
|
| def __init__(self): |
| self.type = "output" |
|
|
| def save(self, mesh, filename_prefix): |
| output_dir = folder_paths.get_output_directory() |
| glbs = [] |
| for idx, m in enumerate(mesh): |
| scene = trimesh.Scene(m) |
| glb_data = gltf.export_glb(scene, include_normals=True) |
| logging.info(f"Generated GLB model with {len(glb_data)} bytes") |
|
|
| full_output_folder, filename, counter, subfolder, filename_prefix = ( |
| folder_paths.get_save_image_path(filename_prefix, output_dir) |
| ) |
| filename = filename.replace("%batch_num%", str(idx)) |
| out_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.glb") |
| with open(out_path, "wb") as f: |
| f.write(glb_data) |
| glbs.append(base64.b64encode(glb_data).decode("utf-8")) |
| return {"ui": {"glbs": glbs}} |
|
|
|
|
| NODE_DISPLAY_NAME_MAPPINGS = { |
| "StableFast3DLoader": "Stable Fast 3D Loader", |
| "StableFast3DPreview": "Stable Fast 3D Preview", |
| "StableFast3DSampler": "Stable Fast 3D Sampler", |
| "StableFast3DSave": "Stable Fast 3D Save", |
| } |
|
|
| NODE_CLASS_MAPPINGS = { |
| "StableFast3DLoader": StableFast3DLoader, |
| "StableFast3DPreview": StableFast3DPreview, |
| "StableFast3DSampler": StableFast3DSampler, |
| "StableFast3DSave": StableFast3DSave, |
| } |
|
|
| WEB_DIRECTORY = "./comfyui" |
|
|
| __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] |
|
|
