DreamWaltzG Stage One Trainer

This node is stage one of a two-stage pipeline for turning a text prompt into a full-body 3D Gaussian avatar you can pose and animate. It wraps DreamWaltz-G (TPAMI 2025, arXiv 2409.17145), a research project that builds an expressive SMPL-X-rigged avatar from scratch using score distillation against Stable Diffusion 1.5. Stage one trains a NeRF initialized from a human template; stage two converts that NeRF into 3D Gaussians. If that sounds heavy, it is - the upstream authors say a full run takes hours on a single L40S.

Here's the honest part you won't get from the pack's README, which literally still says "Features (Planned)": as shipped, this node doesn't actually start the training. In nodes.py the whole python main.py --stage nerf invocation is commented out. What the node really does is set up the experiment folder, point you at the right checkpoint, and hand the path to Stage Two. The actual training call is the one the author left on the floor.

What it actually does

Give it a prompt and it sanitizes it into an experiment name - a chef dressed in white becomes outputs/a_chef_dressed_in_white/nerf_64_64_10k/ under the pack's root folder (the env var DREAMWALTZ_ROOT_FOLDER, printed to your console on load). It creates the checkpoints directory, grabs the newest *.pth there to resume, and falls back to the template checkpoint external/human_templates/instant-ngp/adult_neutral/step_005000.pth if it's a fresh run. It then returns the experiment path plus any preview images already sitting in results/1024x1024/image. So if the folder already has a partially-trained NeRF, this is a resume-and-point node; if it's empty, you've got setup work to do yourself.

Inputs and outputs that matter

• prompt - the text describing your avatar. Be specific; this is your whole dataset.
• iterations - default 10000 (range 1000–100000). NeRF stage cost scales directly with this.
• train_resolution - 64,64 (default), 512,512, or 64,128,256. The three-value option is the progressive-grid scheme: it upsamples as training goes, which is why enable_progressive_grid defaults on.
• background_mode - gray/white/black background for renders. gray is the default and the classic SDS choice.
• body_parts - hands, face, or full. This feeds --predefined_body_parts, i.e. how much of the SMPL-X body gets refined.
• scene_type - canonical or dynamic; canonical (a single neutral pose) is the default for avatar creation.
• enable_fp16 / enable_progressive_grid / enable_sigma_guidance - toggles for mixed precision, the progressive grid, and sigma guidance. Leave them on unless something breaks.

Outputs: checkpoints_folder_path (STRING) is the one that matters - feed it straight into Stage Two's stage_one_checkpoint_path. images (IMAGE) is a preview tensor of whatever renders already exist, and will be None on a fresh run, which is fine.

Installing it

ComfyUI Manager can find it if you search "ComfyUI-DreamWaltz-G", or:

cd ComfyUI/custom_nodes
git clone https://github.com/gt732/ComfyUI-DreamWaltz-G

Then restart ComfyUI. Here's the trap: the pack's install.py is empty, so Manager's dependency auto-install will do nothing. This wrapper assumes you've already set up the upstream DreamWaltz-G environment - a conda env with Python 3.11, torch 2.1.0, pytorch3d 0.7.5, nvdiffrast, smplx, human_body_prior, diff-gaussian-rasterization and controlnet-aux, per scripts/install.sh. That's a genuinely painful stack to build; pip install pytorch3d alone will test your patience. On top of that you need the model files: the SMPL-X/FLAME/VPoser human templates (from the official MPG pages), the pre-trained Instant-NGP template step_005000.pth (HuggingFace KevinHuang/DreamWaltz-G), and Stable Diffusion 1.5, which diffusers pulls from stable-diffusion-v1-5/stable-diffusion-v1-5. Missing any of these, and the run dies on file-not-found before any training happens.

The gotchas

• Nothing trains. Re-check nodes.py before you queue this up - the subprocess call is commented out. Either run the equivalent python main.py --stage nerf ... yourself from the DreamWaltz-G clone, or use this node purely as the folder-and-path bookkeeper it currently is.
• Instant-NGP CUDA extensions build at runtime on first use (they're borrowed from stable-dreamfusion). Expect the first invocation to spend a while compiling.
• It's a VRAM and time sink. Score distillation at 512² is not a 10-minute job; a 10k-iteration NeRF stage on a consumer card will run for hours. Don't babysit it.
• If the node returns instantly with no images and an empty checkpoint dir, that's the commented-out call - not a crash.

This is a sharp-edged early wrapper around serious research code. Fine for someone who already has the upstream environment running and wants to drive it from the graph. If you're hoping for a turnkey "type prompt, get avatar" node, this isn't it yet.