init

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/infinigen"]
+	path = third_party/infinigen
+	url = https://github.com/princeton-vl/infinigen

README.md

@@ -5,6 +5,7 @@ colorFrom: purple
 colorTo: blue
 sdk: gradio
 sdk_version: 5.9.1
+python_version: 3.10.14
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -0,0 +1,258 @@
+import os
+os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
+import yaml
+import numpy as np
+from PIL import Image
+import rembg
+import importlib
+import torch
+import tempfile
+import json
+#import spaces
+from core.models import DiT_models
+from core.diffusion import create_diffusion
+from core.utils.dinov2 import Dinov2Model
+from core.utils.math_utils import unnormalize_params
+
+from huggingface_hub import hf_hub_download
+
+# Setup PyTorch:
+device = torch.device('cuda')
+
+# Define the cache directory for model files
+#model_cache_dir = './ckpts/'
+#os.makedirs(model_cache_dir, exist_ok=True)
+
+# load generators & models
+generators_choices = ["chair", "table", "vase", "basket", "flower", "dandelion"]
+factory_names = ["ChairFactory", "TableDiningFactory", "VaseFactory", "BasketBaseFactory", "FlowerFactory", "DandelionFactory"]
+generator_path = "./core/assets/"
+generators, configs, models = [], [], []
+for category, factory in zip(generators_choices, factory_names):
+    # load generator
+    module = importlib.import_module(f"core.assets.{category}")
+    gen = getattr(module, factory)
+    generator = gen(0)
+    generators.append(generator)
+    # load configs
+    config_path = f"./configs/demo/{category}_demo.yaml"
+    with open(config_path) as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)
+    configs.append(cfg)
+    # load models
+    latent_size = cfg["num_params"]
+    model = DiT_models[cfg["model"]](input_size=latent_size).to(device)
+    # load a custom DiT checkpoint from train.py:
+    # download the checkpoint if not found:
+    if not os.path.exists(cfg["ckpt_path"]):
+        model_dir, model_name = os.path.dirname(cfg["ckpt_path"]), os.path.basename(cfg["ckpt_path"])
+        os.makedirs(model_dir, exist_ok=True)
+        checkpoint_path = hf_hub_download(repo_id="TencentARC/DI-PCG", 
+                            local_dir=model_dir, filename=model_name)
+        print("Downloading checkpoint {} from Hugging Face Hub...".format(model_name))
+    print("Loading model from {}".format(cfg["ckpt_path"]))
+    
+    state_dict = torch.load(cfg["ckpt_path"], map_location=lambda storage, loc: storage)
+    if "ema" in state_dict:  # supports checkpoints from train.py
+        state_dict = state_dict["ema"]
+    model.load_state_dict(state_dict)
+    model.eval()
+    models.append(model)
+    
+diffusion = create_diffusion(str(cfg["num_sampling_steps"]))
+# feature model
+feature_model = Dinov2Model()
+
+
+def check_input_image(input_image):
+    if input_image is None:
+        raise gr.Error("No image uploaded!")
+
+
+def preprocess(input_image, do_remove_background):
+    # resize
+    if input_image.size[0] != 256 or input_image.size[1] != 256:
+        input_image = input_image.resize((256, 256))
+    # remove background
+    if do_remove_background:
+        processed_image = rembg.remove(np.array(input_image))
+    # white background
+    else:
+        processed_image = input_image
+    return processed_image
+
+#@spaces.GPU
+def sample(image, seed, category):
+    # seed
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    # generator & model
+    idx = generators_choices.index(category)
+    generator, cfg, model = generators[idx], configs[idx], models[idx]
+    
+    # encode condition image feature
+    # convert RGBA images to RGB, white background
+    input_image_np = np.array(image)
+    mask = input_image_np[:, :, -1:] > 0
+    input_image_np = input_image_np[:, :, :3] * mask + 255 * (1 - mask)
+    image = input_image_np.astype(np.uint8)
+    
+    img_feat = feature_model.encode_batch_imgs([np.array(image)], global_feat=False)
+
+    # Create sampling noise:
+    latent_size = int(cfg['num_params'])
+    z = torch.randn(1, 1, latent_size, device=device)
+    y = img_feat
+
+    # No classifier-free guidance:
+    model_kwargs = dict(y=y)
+
+    # Sample target params:
+    samples = diffusion.p_sample_loop(
+        model.forward, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+    )
+    samples = samples[0].squeeze(0).cpu().numpy()
+
+    # unnormalize params
+    params_dict = generator.params_dict
+    params_original = unnormalize_params(samples, params_dict)
+    
+    mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False).name
+    params_fpath = tempfile.NamedTemporaryFile(suffix=f".npy", delete=False).name
+    np.save(params_fpath, params_original)
+    print(mesh_fpath)
+    print(params_fpath)
+    # generate 3D using sampled params - TODO: this is a hacky way to go through PCG pipeline, avoiding conflict with gradio
+    command = f"python ./scripts/generate.py --config ./configs/demo/{category}_demo.yaml --output_path {mesh_fpath} --seed {seed} --params_path {params_fpath}"
+    os.system(command)
+    
+    return mesh_fpath
+
+
+import gradio as gr
+
+_HEADER_ = '''
+<h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/TencentARC/DI-PCG' target='_blank'><b>DI-PCG: Diffusion-based Efficient Inverse Procedural Content Generation for High-quality 3D Asset Creation</b></a></h2>
+
+**DI-PCG** is a diffusion model which directly generates a procedural generator's parameters from a single image, resulting in high-quality 3D meshes.
+
+Code: <a href='https://github.com/TencentARC/DI-PCG' target='_blank'>GitHub</a>. Techenical report: <a href='' target='_blank'>ArXiv</a>.
+
+❗️❗️❗️**Important Notes:**
+- DI-PCG trains a diffusion model for each procedural generator. Current supported generators are: Chair, Table, Vase, Basket, Flower, Dandelion from <a href="https://github.com/princeton-vl/infinigen">Infinigen</a>.
+- The diversity of the generated meshes are strictly bounded by the procedural generators. For out-of-domain shapes, DI-PCG may only provide closest approximations.
+'''
+
+_CITE_ = r"""
+If DI-PCG is helpful, please help to ⭐ the <a href='https://github.com/TencentARC/DI-PCG' target='_blank'>Github Repo</a>. Thanks! [![GitHub Stars](https://img.shields.io/github/stars/TencentARC/DI-PCG?style=social)](https://github.com/TencentARC/DI-PCG)
+---
+📝 **Citation**
+
+If you find our work useful for your research or applications, please cite using this bibtex:
+```bibtex
+
+```
+
+📋 **License**
+
+Apache-2.0 LICENSE. Please refer to the [LICENSE file]() for details.
+
+📧 **Contact**
+
+If you have any questions, feel free to open a discussion or contact us at <b></b>.
+"""
+def update_examples(category):
+    samples = [[os.path.join(f"examples/{category}", img_name)]
+        for img_name in sorted(os.listdir(f"examples/{category}"))]
+    print(samples)
+    return gr.Dataset(samples=samples)
+
+with gr.Blocks() as demo:
+    gr.Markdown(_HEADER_)
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            # select the generator category
+            with gr.Row():
+                with gr.Group():
+                    generator_category = gr.Radio(
+                        choices=[
+                            "chair",
+                            "table",
+                            "vase",
+                            "basket",
+                            "flower",
+                            "dandelion",
+                        ],
+                        value="chair",
+                        label="category",
+                    )
+            with gr.Row():
+                input_image = gr.Image(
+                    label="Input Image",
+                    image_mode="RGB",
+                    sources='upload',
+                    width=256,
+                    height=256,
+                    type="pil",
+                    elem_id="content_image",
+                )
+                processed_image = gr.Image(
+                    label="Processed Image", 
+                    image_mode="RGBA", 
+                    width=256,
+                    height=256,
+                    type="pil", 
+                    interactive=False
+                )
+            with gr.Row():
+                with gr.Group():
+                    do_remove_background = gr.Checkbox(
+                        label="Remove Background", value=False
+                    )
+                    sample_seed = gr.Number(value=0, label="Seed Value", precision=0)
+
+            with gr.Row():
+                submit = gr.Button("Generate", elem_id="generate", variant="primary")
+
+            with gr.Row(variant="panel"):
+                examples = gr.Examples(
+                    [os.path.join(f"examples/chair", img_name) for img_name in sorted(os.listdir(f"examples/chair"))],
+                    inputs=[input_image],
+                    label="Examples",
+                    examples_per_page=5
+                )
+                generator_category.change(update_examples, generator_category, outputs=examples.dataset)
+
+        with gr.Column():
+            with gr.Row():
+                with gr.Tab("Geometry"):
+                    output_model_obj = gr.Model3D(
+                        label="Output Model",
+                        #width=768,
+                        display_mode="wireframe",
+                        interactive=False
+                    )
+                #with gr.Tab("Textured"):
+                #    output_model_obj = gr.Model3D(
+                #        label="Output Model (STL Format)",
+                #        #width=768,
+                #        interactive=False,
+                #    )
+                #    gr.Markdown("Note: Texture and Material are randomly assigned by the procedural generator.")
+
+
+    gr.Markdown(_CITE_)
+    mv_images = gr.State()
+
+    submit.click(fn=check_input_image, inputs=[input_image]).success(
+        fn=preprocess,
+        inputs=[input_image, do_remove_background],
+        outputs=[processed_image],
+    ).success(
+        fn=sample,
+        inputs=[processed_image, sample_seed, generator_category],
+        outputs=[output_model_obj],
+    )
+
+demo.queue(max_size=10)
+demo.launch(server_name="0.0.0.0", server_port=43839)

configs/demo/basket_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/basket
+save_dir: logs/basket_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/basket.pt
+
+# Generator
+generator_root: core/assets
+generator: BasketBaseFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 14
+
+# Render
+r_cam_dists: [1.6]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.0

configs/demo/chair_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/chair
+save_dir: logs/chair_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/chair.pt
+
+# Generator
+generator_root: core/assets
+generator: ChairFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 48
+
+# Render
+r_cam_dists: [2.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.0

configs/demo/dandelion_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/dandelion
+save_dir: logs/dandelion_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/dandelion.pt
+
+# Generator
+generator_root: core/assets
+generator: DandelionFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 15
+
+# Render
+r_cam_dists: [3.0]
+r_cam_elevations: [90]
+r_cam_azimuths: [0]
+r_zoff: 0.5

configs/demo/flower_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/flower
+save_dir: logs/flower_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/flower.pt
+
+# Generator
+generator_root: core/assets
+generator: FlowerFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 9
+
+# Render
+r_cam_dists: [4.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [0]
+r_zoff: 0.0

configs/demo/table_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/table
+save_dir: logs/table_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/table.pt
+
+# Generator
+generator_root: core/assets
+generator: TableDiningFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 19
+
+# Render
+r_cam_dists: [5.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.1

configs/demo/vase_demo.yaml

@@ -0,0 +1,20 @@
+# Test
+condition_img_dir: examples/vase
+save_dir: logs/vase_demo
+num_sampling_steps: 250
+ckpt_path: pretrained_models/vase.pt
+
+# Generator
+generator_root: core/assets
+generator: VaseFactory
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 12
+
+# Render
+r_cam_dists: [2.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [0]
+r_zoff: 0.3

configs/infinigen/base.gin

@@ -0,0 +1,89 @@
+include 'surface_registry.gin'
+
+OVERALL_SEED = 0
+LOG_DIR = '.'
+
+Terrain.asset_folder = "" # Will read from $INFINIGEN_ASSET_FOLDER environment var when set to None, and on the fly when set to ""
+Terrain.asset_version = 'May27'
+
+util.math.FixedSeed.seed = %OVERALL_SEED
+
+execute_tasks.frame_range = [1, 1] # Between start/end frames should this job consider? Increase end frame to tackle video
+execute_tasks.camera_id = [0, 0] # Which camera rig
+
+save_obj_and_instances.output_folder="saved_mesh.obj"
+
+util.logging.create_text_file.log_dir = %LOG_DIR
+
+target_face_size.global_multiplier = 2
+scatter_res_distance.dist = 4
+
+random_color_mapping.hue_stddev = 0.05 # Note: 1.0 is the whole color spectrum
+
+render.render_image_func = @full/render_image
+configure_render_cycles.time_limit = 0
+
+configure_render_cycles.min_samples = 0
+configure_render_cycles.num_samples = 8192
+configure_render_cycles.adaptive_threshold = 0.01
+configure_render_cycles.denoise = False
+configure_render_cycles.exposure = 1
+configure_blender.motion_blur_shutter = 0.15
+render_image.use_dof = False
+render_image.dof_aperture_fstop = 3
+compositor_postprocessing.distort = False
+compositor_postprocessing.color_correct = False
+
+flat/configure_render_cycles.min_samples = 1
+flat/configure_render_cycles.num_samples = 16
+flat/render_image.flat_shading = True
+full/render_image.passes_to_save = [
+    ['diffuse_direct', 'DiffDir'],
+    ['diffuse_color', 'DiffCol'],
+    ['diffuse_indirect', 'DiffInd'],
+    ['glossy_direct', 'GlossDir'],
+    ['glossy_color', 'GlossCol'],
+    ['glossy_indirect', 'GlossInd'],
+    ['transmission_direct', 'TransDir'],
+    ['transmission_color', 'TransCol'],
+    ['transmission_indirect', 'TransInd'],
+    ['volume_direct', 'VolumeDir'],
+    ['emit', 'Emit'],
+    ['environment', 'Env'],
+    ['ambient_occlusion', 'AO']
+]
+flat/render_image.passes_to_save = [
+    ['z', 'Depth'],
+    ['normal', 'Normal'],
+    ['vector', 'Vector'],
+    ['object_index', 'IndexOB']
+]
+
+execute_tasks.generate_resolution = (1280, 720)
+execute_tasks.fps = 24
+get_sensor_coords.H = 720
+get_sensor_coords.W = 1280
+
+min_terrain_distance = 2
+keep_cam_pose_proposal.min_terrain_distance = %min_terrain_distance
+SphericalMesher.r_min = %min_terrain_distance
+
+build_terrain_bvh_and_attrs.avoid_border = False # disabled due to crashes 5/15
+
+animate_cameras.follow_poi_chance=0.0
+camera.camera_pose_proposal.altitude = ("weighted_choice",
+    (0.975, ("clip_gaussian", 2, 0.3, 0.5, 3)), # person height usually
+    (0.025, ("clip_gaussian", 15, 7, 5, 30)) # drone height sometimes
+)
+
+camera.camera_pose_proposal.pitch = ("clip_gaussian", 90, 30, 20, 160)
+
+# WARNING: Large camera rig translations or rotations require special handling. 
+#    if your cameras are not all approximately forward facing within a few centimeters, you must either:
+#    -  configure the pipeline to generate assets / terrain for each camera separately, rather than sharing it between the whole rig
+#    -  or, treat your camera rig as multiple camera rigs each with one camera, and implement code to positon them correctly
+camera.spawn_camera_rigs.n_camera_rigs = 1
+camera.spawn_camera_rigs.camera_rig_config = [
+    {'loc': (0, 0, 0), 'rot_euler': (0, 0, 0)},
+    {'loc': (0.075, 0, 0), 'rot_euler': (0, 0, 0)}
+]

configs/test/basket_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/basket_test
+data_root: /group/40034/wangzhao/data/ipcg/basket_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+ 
+# Generator
+run_generate: False
+generator: BasketBaseFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 14
+
+# Render
+r_cam_dists: [1.6]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.0

configs/test/chair_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/chair_test
+data_root: /group/40034/wangzhao/data/ipcg/chair_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+ 
+# Generator
+run_generate: True
+generator: ChairFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 48
+
+# Render
+r_cam_dists: [2.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.0

configs/test/dandelion_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/dandelion_test
+data_root: /group/40075/wangzhao/ipcg/dandelion_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+
+# Generator
+run_generate: True
+generator: DandelionFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 15
+
+# Render
+r_cam_dists: [3.0]
+r_cam_elevations: [90]
+r_cam_azimuths: [0]
+r_zoff: 0.5

configs/test/flower_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/flower_test
+data_root: /group/40075/wangzhao/ipcg/flower_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+
+# Generator
+run_generate: True 
+generator: FlowerFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 9
+
+# Render
+r_cam_dists: [4.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [0]
+r_zoff: 0.0

configs/test/table_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/table_test
+data_root: /group/40034/wangzhao/data/ipcg/table_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+
+# Generator
+run_generate: True 
+generator: TableDiningFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 19
+
+# Render
+r_cam_dists: [5.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [30]
+r_zoff: 0.1

configs/test/vase_test.yaml

@@ -0,0 +1,24 @@
+# Test
+save_dir: logs/vase_test
+data_root: /group/40034/wangzhao/data/ipcg/vase_new
+test_file: test_list_mv.txt
+batch_size: 100
+num_workers: 24
+num_sampling_steps: 250
+ckpt_path: /your/path/to/trained/model/ckpt.pt
+
+# Generator
+run_generate: True
+generator: VaseFactory
+params_dict_file: params_dict.txt
+seed: 0
+
+# Model
+model: DiT_mini
+num_params: 12
+
+# Render
+r_cam_dists: [2.0]
+r_cam_elevations: [60]
+r_cam_azimuths: [0]
+r_zoff: 0.3

configs/train/basket_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/basket_train
+data_root: /group/40075/wangzhao/ipcg/basket
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 14

configs/train/chair_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/chair_train
+data_root: /group/40046/public_datasets/IPCG/chair_new
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 48

configs/train/dandelion_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/dandelion_train
+data_root: /group/40034/wangzhao/data/ipcg/dandelion
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 15

configs/train/flower_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/flower_train
+data_root: /workspace/40075_wangzhao/ipcg/flower_new
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 9

configs/train/table_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/table_train
+data_root: /group/40034/wangzhao/data/ipcg/table_new
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 19

configs/train/vase_train.yaml

@@ -0,0 +1,17 @@
+# Train
+save_dir: logs/vase_train
+data_root: /group/40034/wangzhao/data/ipcg/vase_new
+train_file: train_list_mv_withaug.txt
+test_file: test_list_mv.txt
+params_dict_file: params_dict.txt
+epochs: 200
+batch_size: 128
+num_workers: 64
+lr: 0.0001
+seed: 0
+logging_iter: 100
+ckpt_iter: 10000
+
+# Model
+model: DiT_mini
+num_params: 12

core/assets/basket.py

@@ -0,0 +1,576 @@
+# Copyright (C) 2023, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Beining Han
+
+import bpy
+import numpy as np
+from numpy.random import uniform
+import random
+import time
+
+from infinigen.assets.materials.plastics.plastic_rough import shader_rough_plastic
+from infinigen.core import surface, tagging
+from infinigen.core.nodes import node_utils
+from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
+from infinigen.core.placement.factory import AssetFactory
+
+
+@node_utils.to_nodegroup("nodegroup_holes", singleton=False, type="GeometryNodeTree")
+def nodegroup_holes(nw: NodeWrangler):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketFloat", "height", 0.5000),
+            ("NodeSocketFloat", "gap_size", 0.5000),
+            ("NodeSocketFloat", "hole_edge_gap", 0.5000),
+            ("NodeSocketFloat", "hole_size", 0.5000),
+            ("NodeSocketFloat", "depth", 0.5000),
+            ("NodeSocketFloat", "width", 0.5000),
+        ],
+    )
+
+    add = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["hole_edge_gap"], 1: 0.0000},
+        attrs={"operation": "ADD"}
+    )
+
+    subtract = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["height"], 1: add},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    add_1 = nw.new_node(
+        Nodes.Math, 
+        input_kwargs={0: group_input.outputs["width"], 1: 0.0000},
+        attrs={"operation": "ADD"}
+    )
+
+    subtract_1 = nw.new_node(
+        Nodes.Math, input_kwargs={0: add_1, 1: add}, attrs={"operation": "SUBTRACT"}
+    )
+
+    add_2 = nw.new_node(
+        Nodes.Math, input_kwargs={0: group_input.outputs["hole_size"], 1: 0.0000}, attrs={"operation": "ADD"}
+    )
+
+    add_3 = nw.new_node(
+        Nodes.Math, input_kwargs={0: add_2, 1: group_input.outputs["gap_size"]}, attrs={"operation": "ADD"}
+    )
+
+    divide = nw.new_node(
+        Nodes.Math, input_kwargs={0: subtract, 1: add_3}, attrs={"operation": "DIVIDE"}
+    )
+
+    divide_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: subtract_1, 1: add_3},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    grid = nw.new_node(
+        Nodes.MeshGrid,
+        input_kwargs={
+            "Size X": subtract,
+            "Size Y": subtract_1,
+            "Vertices X": divide,
+            "Vertices Y": divide_1,
+        },
+    )
+
+    store_named_attribute = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": grid.outputs["Mesh"],
+            "Name": "uv_map",
+            3: grid.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    transform_1 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": store_named_attribute,
+            "Rotation": (0.0000, 1.5708, 0.0000),
+        },
+    )
+
+    add_4 = nw.new_node(
+        Nodes.Math, input_kwargs={0: group_input.outputs["depth"], 1: 0.0000}, attrs={"operation": "ADD"}
+    )
+
+    add_5 = nw.new_node(Nodes.Math, input_kwargs={0: add_4, 1: 0.1}, attrs={"operation": "ADD"})
+
+    combine_xyz_3 = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": add_5, "Y": add_2, "Z": add_2}
+    )
+
+    cube_2 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_3})
+
+    store_named_attribute_1 = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": cube_2.outputs["Mesh"],
+            "Name": "uv_map",
+            3: cube_2.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    instance_on_points = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={"Points": transform_1, "Instance": store_named_attribute_1},
+    )
+
+    subtract_2 = nw.new_node(
+        Nodes.Math, input_kwargs={0: add_4, 1: add}, attrs={"operation": "SUBTRACT"}
+    )
+
+    divide_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: subtract_2, 1: add_3},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    grid_1 = nw.new_node(
+        Nodes.MeshGrid,
+        input_kwargs={
+            "Size X": subtract_2,
+            "Size Y": subtract,
+            "Vertices X": divide_2,
+            "Vertices Y": divide,
+        },
+    )
+
+    store_named_attribute_2 = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": grid_1.outputs["Mesh"],
+            "Name": "uv_map",
+            3: grid_1.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    transform_2 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": store_named_attribute_2,
+            "Rotation": (1.5708, 0.0000, 0.0000),
+        },
+    )
+
+    add_6 = nw.new_node(Nodes.Math, input_kwargs={0: add_1, 1: 0.1}, attrs={"operation": "ADD"})
+
+    combine_xyz_4 = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": add_2, "Y": add_6, "Z": add_2}
+    )
+
+    cube_3 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_4})
+
+    store_named_attribute_3 = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": cube_3.outputs["Mesh"],
+            "Name": "uv_map",
+            3: cube_3.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    instance_on_points_1 = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={"Points": transform_2, "Instance": store_named_attribute_3},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={
+            "Instances1": instance_on_points,
+            "Instances2": instance_on_points_1,
+        },
+        attrs={"is_active_output": True},
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_handle_hole", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_handle_hole(nw: NodeWrangler):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketFloat", "X", 0.0000),
+            ("NodeSocketFloat", "Z", 0.0000),
+            ("NodeSocketFloat", "height", 0.5000),
+            ("NodeSocketFloat", "hole_dist", 0.5000),
+            ("NodeSocketInt", "Level", 0),
+        ],
+    )
+
+    combine_xyz_3 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={
+            "X": group_input.outputs["X"],
+            "Y": 1.0000,
+            "Z": group_input.outputs["Z"],
+        },
+    )
+
+    cube_2 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_3})
+
+    store_named_attribute = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": cube_2.outputs["Mesh"],
+            "Name": "uv_map",
+            3: cube_2.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    subdivide_mesh_2 = nw.new_node(
+        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute}
+    )
+
+    subdivision_surface_2 = nw.new_node(
+        Nodes.SubdivisionSurface,
+        input_kwargs={"Mesh": subdivide_mesh_2, "Level": group_input.outputs["Level"]},
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["height"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    subtract = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: multiply, 1: group_input.outputs["hole_dist"]},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    combine_xyz_4 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": subtract})
+
+    transform_1 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={"Geometry": subdivision_surface_2, "Translation": combine_xyz_4},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": transform_1},
+        attrs={"is_active_output": True},
+    )
+
+
+def geometry_nodes(nw: NodeWrangler, **kwargs):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    depth = nw.new_node(Nodes.Value, label="depth")
+    depth.outputs[0].default_value = kwargs["depth"]
+
+    width = nw.new_node(Nodes.Value, label="width")
+    width.outputs[0].default_value = kwargs["width"]
+
+    height = nw.new_node(Nodes.Value, label="height")
+    height.outputs[0].default_value = kwargs["height"]
+
+    combine_xyz = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": depth, "Y": width, "Z": height}
+    )
+
+    cube = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz})
+
+    store_named_attribute = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": cube.outputs["Mesh"],
+            "Name": "uv_map",
+            3: cube.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    subdivide_mesh = nw.new_node(
+        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute, "Level": 2}
+    )
+
+    sub_level = nw.new_node(Nodes.Integer, label="sub_level")
+    sub_level.integer = kwargs["frame_sub_level"]
+
+    subdivision_surface = nw.new_node(
+        Nodes.SubdivisionSurface,
+        input_kwargs={"Mesh": subdivide_mesh, "Level": sub_level},
+    )
+
+    differences = []
+
+    if kwargs["has_handle"]:
+        hole_depth = nw.new_node(Nodes.Value, label="hole_depth")
+        hole_depth.outputs[0].default_value = kwargs["handle_depth"]
+
+        hole_height = nw.new_node(Nodes.Value, label="hole_height")
+        hole_height.outputs[0].default_value = kwargs["handle_height"]
+
+        hole_dist = nw.new_node(Nodes.Value, label="hole_dist")
+        hole_dist.outputs[0].default_value = kwargs["handle_dist_to_top"]
+
+        handle_level = nw.new_node(Nodes.Integer, label="handle_level")
+        handle_level.integer = kwargs["handle_sub_level"]
+        handle_hole = nw.new_node(
+            nodegroup_handle_hole().name,
+            input_kwargs={
+                "X": hole_depth,
+                "Z": hole_height,
+                "height": height,
+                "hole_dist": hole_dist,
+                "Level": handle_level,
+            },
+        )
+        differences.append(handle_hole)
+
+    thickness = nw.new_node(Nodes.Value, label="thickness")
+    thickness.outputs[0].default_value = kwargs["thickness"]
+
+    subtract = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: depth, 1: thickness},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    subtract_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: width, 1: thickness},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": subtract, "Y": subtract_1, "Z": height}
+    )
+
+    cube_1 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_1})
+
+    store_named_attribute_1 = nw.new_node(
+        Nodes.StoreNamedAttribute,
+        input_kwargs={
+            "Geometry": cube_1.outputs["Mesh"],
+            "Name": "uv_map",
+            3: cube_1.outputs["UV Map"],
+        },
+        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
+    )
+
+    subdivide_mesh_1 = nw.new_node(
+        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute_1, "Level": 2}
+    )
+
+    subdivision_surface_1 = nw.new_node(
+        Nodes.SubdivisionSurface,
+        input_kwargs={"Mesh": subdivide_mesh_1, "Level": sub_level},
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: thickness, 2: 0.2500},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_2 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply})
+
+    transform = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={"Geometry": subdivision_surface_1, "Translation": combine_xyz_2},
+    )
+
+    if kwargs["has_holes"]:
+        gap_size = nw.new_node(Nodes.Value, label="gap_size")
+        gap_size.outputs[0].default_value = kwargs["hole_gap_size"]
+
+        hole_edge_gap = nw.new_node(Nodes.Value, label="hole_edge_gap")
+        hole_edge_gap.outputs[0].default_value = kwargs["hole_edge_gap"]
+
+        hole_size = nw.new_node(Nodes.Value, label="hole_size")
+        hole_size.outputs[0].default_value = kwargs["hole_size"]
+        holes = nw.new_node(
+            nodegroup_holes().name,
+            input_kwargs={
+                "height": height,
+                "gap_size": gap_size,
+                "hole_edge_gap": hole_edge_gap,
+                "hole_size": hole_size,
+                "depth": depth,
+                "width": width,
+            },
+        )
+        differences.extend([holes.outputs["Instances1"], holes.outputs["Instances2"]])
+
+    difference = nw.new_node(
+        Nodes.MeshBoolean,
+        input_kwargs={
+            "Mesh 1": subdivision_surface,
+            "Mesh 2": [transform] + differences,
+        },
+    )
+
+    realize_instances = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": difference.outputs["Mesh"]}
+    )
+
+    multiply_1 = nw.new_node(
+        Nodes.Math, input_kwargs={0: height}, attrs={"operation": "MULTIPLY"}
+    )
+
+    combine_xyz_3 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_1})
+
+    transform_geometry = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={"Geometry": realize_instances, "Translation": combine_xyz_3},
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": transform_geometry,
+            "Material": surface.shaderfunc_to_material(shader_rough_plastic),
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": set_material},
+        attrs={"is_active_output": True},
+    )
+
+
+class BasketBaseFactory(AssetFactory):
+    def __init__(self, factory_seed, coarse=False):
+        super(BasketBaseFactory, self).__init__(factory_seed, coarse=coarse)
+        self.params = self.get_asset_params()
+        self.seed = factory_seed
+        self.get_params_dict()
+    
+    def get_params_dict(self):
+        self.params_dict = {
+            "depth": ['continuous', (0.1, 0.6)],
+            "width": ['continuous', (0.1, 0.7)],
+            "height": ['continuous', (0.05, 0.4)],
+            "frame_sub_level": ['discrete', [0, 3]],
+            "thickness": ['continuous', (0.001, 0.03)],
+            "has_handle": ['discrete', [0, 1]],
+            "handle_sub_level": ['discrete', [0, 1, 2]],
+            "handle_depth": ['continuous', (0.2, 0.6)],
+            "handle_height": ['continuous', (0.1, 0.3)],
+            "handle_dist_to_top": ['continuous', (0.08, 0.4)],
+            "has_holes": ['discrete', [0, 1]],
+            "hole_gap_size": ['continuous', (0.5, 2.0)],
+            "hole_edge_gap": ['continuous', (0.04, 0.1)],
+            "hole_size": ['continuous', (0.007, 0.02)]
+        }
+
+    def fix_unused_params(self, params):
+        if params['height'] < 0.12:
+            params['has_holes'] = 0
+        if params['has_handle'] == 0:
+            params["handle_sub_level"] = 1
+            params["handle_depth"] = 0.3
+            params["handle_height"] = 0.2
+            params["handle_dist_to_top"] = 0.115
+        if params['has_holes'] == 0:
+            params["hole_gap_size"] = 0.95
+            params["hole_edge_gap"] = 0.05
+            params["hole_size"] = 0.0075
+        return params
+
+    def update_params(self, params):
+        # TODO: to allow random material
+        self.seed = int(1000 * time.time()) % 2**32
+        
+        handle_depth = params['depth'] * params['handle_depth']
+        handle_height = params['height'] * params['handle_height']
+        handle_dist_to_top = handle_height * 0.5 + params['height'] * params["handle_dist_to_top"]
+        if params['height'] < 0.12:
+            params["has_holes"] = 0
+        hole_gap_size = params['hole_size'] * params["hole_gap_size"]
+        parameters = {
+            "depth": params["depth"],
+            "width": params["width"],
+            "height": params["height"],
+            "frame_sub_level": params["frame_sub_level"],
+            "thickness": params["thickness"],
+            "has_handle": params["has_handle"] > 0,
+            "handle_sub_level": params["handle_sub_level"],
+            "handle_depth": handle_depth,
+            "handle_height": handle_height,
+            "handle_dist_to_top": handle_dist_to_top,
+            "has_holes": params["has_holes"] > 0,
+            "hole_gap_size": hole_gap_size,
+            "hole_edge_gap": params["hole_edge_gap"],
+            "hole_size": params["hole_size"],
+        }
+        self.params.update(parameters)
+
+    def get_asset_params(self, i=0):
+        params = {}
+        if params.get("depth", None) is None:
+            params["depth"] = uniform(0.15, 0.4)
+        if params.get("width", None) is None:
+            params["width"] = uniform(0.2, 0.6)
+        if params.get("height", None) is None:
+            params["height"] = uniform(0.06, 0.24)
+        if params.get("frame_sub_level", None) is None:
+            params["frame_sub_level"] = np.random.choice([0, 3], p=[0.5, 0.5])
+        if params.get("thickness", None) is None:
+            params["thickness"] = uniform(0.001, 0.005)
+
+        if params.get("has_handle", None) is None:
+            params["has_handle"] = np.random.choice([True, False], p=[0.8, 0.2])
+        if params.get("handle_sub_level", None) is None:
+            params["handle_sub_level"] = np.random.choice([0, 1, 2], p=[0.2, 0.4, 0.4])
+        if params.get("handle_depth", None) is None:
+            params["handle_depth"] = params["depth"] * uniform(0.2, 0.4)
+        if params.get("handle_height", None) is None:
+            params["handle_height"] = params["height"] * uniform(0.1, 0.25)
+        if params.get("handle_dist_to_top", None) is None:
+            params["handle_dist_to_top"] = params["handle_height"] * 0.5 + params[
+                "height"
+            ] * uniform(0.08, 0.15)
+
+        if params.get("has_holes", None) is None:
+            if params["height"] < 0.12:
+                params["has_holes"] = False
+            else:
+                params["has_holes"] = np.random.choice([True, False], p=[0.5, 0.5])
+        if params.get("hole_size", None) is None:
+            params["hole_size"] = uniform(0.005, 0.01)
+        if params.get("hole_gap_size", None) is None:
+            params["hole_gap_size"] = params["hole_size"] * uniform(0.8, 1.1)
+        if params.get("hole_edge_gap", None) is None:
+            params["hole_edge_gap"] = uniform(0.04, 0.06)
+
+        return params
+
+    def create_asset(self, i=0, **params):
+        bpy.ops.mesh.primitive_plane_add(
+            size=1,
+            enter_editmode=False,
+            align="WORLD",
+            location=(0, 0, 0),
+            scale=(1, 1, 1),
+        )
+        obj = bpy.context.active_object
+        np.random.seed(self.seed)
+        random.seed(self.seed)
+        
+        surface.add_geomod(
+            obj, geometry_nodes, attributes=[], apply=True, input_kwargs=self.params
+        )
+        tagging.tag_system.relabel_obj(obj)
+        return obj

core/assets/chair.py

@@ -0,0 +1,657 @@
+# Copyright (C) 2024, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Lingjie Mei
+import bpy
+import numpy as np
+from numpy.random import uniform
+
+import infinigen
+from infinigen.assets.material_assignments import AssetList
+from infinigen.assets.utils.decorate import (
+    read_co,
+    read_edge_center,
+    read_edge_direction,
+    remove_edges,
+    remove_vertices,
+    select_edges,
+    solidify,
+    subsurf,
+    write_attribute,
+    write_co,
+)
+from infinigen.assets.utils.draw import align_bezier, bezier_curve
+from infinigen.assets.utils.nodegroup import geo_radius
+from infinigen.assets.utils.object import join_objects, new_bbox
+from infinigen.core import surface
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.surface import NoApply
+from infinigen.core.util import blender as butil
+from infinigen.core.util.blender import deep_clone_obj
+from infinigen.core.util.math import FixedSeed
+from infinigen.core.util.random import log_uniform
+from infinigen.core.util.random import random_general as rg
+
+
+class ChairFactory(AssetFactory):
+    back_types = {
+        0: "whole",
+        1: "partial",
+        2: "horizontal-bar",
+        3: "vertical-bar",
+    }
+    leg_types = {
+        0: "vertical",
+        1: "straight",
+        2: "up-curved",
+        3: "down-curved",
+    }
+
+    def __init__(self, factory_seed, coarse=False):
+        super().__init__(factory_seed, coarse)
+
+        self.get_params_dict()
+        # random init with seed
+        with FixedSeed(self.factory_seed):
+            self.width = uniform(0.4, 0.5)
+            self.size = uniform(0.38, 0.45)
+            self.thickness = uniform(0.04, 0.08)
+            self.bevel_width = self.thickness * (0.1 if uniform() < 0.4 else 0.5)
+            self.seat_back = uniform(0.7, 1.0) if uniform() < 0.75 else 1.0
+            self.seat_mid = uniform(0.7, 0.8)
+            self.seat_mid_x = uniform(
+                self.seat_back + self.seat_mid * (1 - self.seat_back), 1
+            )
+            self.seat_mid_z = uniform(0, 0.5)
+            self.seat_front = uniform(1.0, 1.2)
+            self.is_seat_round = uniform() < 0.6
+            self.is_seat_subsurf = uniform() < 0.5
+
+            self.leg_thickness = uniform(0.04, 0.06)
+            self.limb_profile = uniform(1.5, 2.5)
+            self.leg_height = uniform(0.45, 0.5)
+            self.back_height = uniform(0.4, 0.5)
+            self.is_leg_round = uniform() < 0.5
+            self.leg_type = np.random.choice(
+                ["vertical", "straight", "up-curved", "down-curved"]
+            )
+
+            self.leg_x_offset = 0
+            self.leg_y_offset = 0, 0
+            self.back_x_offset = 0
+            self.back_y_offset = 0
+
+            self.has_leg_x_bar = uniform() < 0.6
+            self.has_leg_y_bar = uniform() < 0.6
+            self.leg_offset_bar = uniform(0.2, 0.4), uniform(0.6, 0.8)
+
+            self.has_arm = uniform() < 0.7
+            self.arm_thickness = uniform(0.04, 0.06)
+            self.arm_height = self.arm_thickness * uniform(0.6, 1)
+            self.arm_y = uniform(0.8, 1) * self.size
+            self.arm_z = uniform(0.3, 0.6) * self.back_height
+            self.arm_mid = np.array(
+                [uniform(-0.03, 0.03), uniform(-0.03, 0.09), uniform(-0.09, 0.03)]
+            )
+            self.arm_profile = log_uniform(0.1, 3, 2)
+
+            self.back_thickness = uniform(0.04, 0.05)
+            self.back_type = rg(self.back_types)
+            self.back_profile = [(0, 1)]
+            self.back_vertical_cuts = np.random.randint(1, 4)
+            self.back_partial_scale = uniform(1, 1.4)
+
+            materials = AssetList["ChairFactory"]()
+            self.limb_surface = materials["limb"].assign_material()
+            self.surface = materials["surface"].assign_material()
+            if uniform() < 0.3:
+                self.panel_surface = self.surface
+            else:
+                self.panel_surface = materials["panel"].assign_material()
+
+            scratch_prob, edge_wear_prob = materials["wear_tear_prob"]
+            self.scratch, self.edge_wear = materials["wear_tear"]
+            is_scratch = uniform() < scratch_prob
+            is_edge_wear = uniform() < edge_wear_prob
+            if not is_scratch:
+                self.scratch = None
+            if not is_edge_wear:
+                self.edge_wear = None
+
+            # from infinigen.assets.clothes import blanket
+            # from infinigen.assets.scatters.clothes import ClothesCover
+            # self.clothes_scatter = ClothesCover(factory_fn=blanket.BlanketFactory, width=log_uniform(.8, 1.2),
+            #                                    size=uniform(.8, 1.2)) if uniform() < .3 else NoApply()
+            self.clothes_scatter = NoApply()
+            self.post_init()
+    
+    def get_params_dict(self):
+        # all the parameters (key:name, value: [type, range]) used in this generator
+        self.params_dict = {
+            "width": ['continuous', [0.3, 0.8]], # seat width
+            "size": ['continuous', [0.35, 0.5]], # seat length
+            "thickness": ['continuous', [0.02, 0.1]], # seat thickness
+            "bevel_width": ['discrete', [0.1, 0.5]],
+            "seat_back": ['continuous', [0.6, 1.0]], # seat back width
+            "seat_mid": ['continuous', [0.7, 0.8]],
+            "seat_mid_z": ['continuous', [0.0, 0.7]], # seat mid point height
+            "seat_front": ['continuous', [1.0, 1.2]], # seat front point
+            "is_seat_round": ['discrete', [0, 1]],
+            "is_seat_subsurf": ['discrete', [0, 1]],
+            "leg_thickness": ['continuous', [0.02, 0.07]], # leg thickness
+            "limb_profile": ['continuous', [1.5, 2.5]],
+            "leg_height": ['continuous', [0.2, 1.0]], # leg height
+            "is_leg_round": ['discrete', [0, 1]],
+            "leg_type": ['discrete', [0,1,2,3]],
+            "has_leg_x_bar": ['discrete', [0, 1]],
+            "has_leg_y_bar": ['discrete', [0, 1]],
+            "leg_offset_bar0": ['continuous', [0.1, 0.9]], # leg y bar offset, only for has_leg_y_bar is 1
+            "leg_offset_bar1": ['continuous', [0.1, 0.9]], # leg x bar offset, only for has_leg_x_bar is 1
+            "leg_x_offset": ['continuous', [0.0, 0.2]], # leg end point x offset
+            "leg_y_offset0": ['continuous', [0.0, 0.2]],  # leg end point y offset
+            "leg_y_offset1": ['continuous', [0.0, 0.2]],  # leg end point y offset
+            "has_arm": ['discrete', [0, 1]],
+            "arm_thickness": ['continuous', [0.02, 0.07]], # arm thickness, only for has_arm is 1
+            "arm_height": ['continuous', [0.6, 1]], # only for has_arm is 1
+            "arm_y": ['continuous', [0.5, 1]], # arm y end point, only for has_arm is 1
+            "arm_z": ['continuous', [0.25, 0.6]], # arm z end point, only for has_arm is 1
+            "arm_mid0": ['continuous', [-0.03, 0.03]], # arm mid point x coord, only for has_arm is 1
+            "arm_mid1": ['continuous', [-0.03, 0.2]], # arm mid point y coord, only for has_arm is 1
+            "arm_mid2": ['continuous', [-0.09, 0.03]], # arm mid point z coord, only for has_arm is 1
+            "arm_profile0": ['continuous', [0.0, 2.0]], # arm curve control, only for has_arm is 1
+            "arm_profile1": ['continuous', [0.0, 2]], # arm curve control, only for has_arm is 1
+            "back_height": ['continuous', [0.3, 0.6]], # back height
+            "back_thickness": ['continuous', [0.02, 0.07]], # back thickness
+            "back_type": ['discrete', [0, 1, 2, 3]],
+            "back_vertical_cuts": ['discrete', [1,2,3,4]], # only for back type 3
+            "back_partial_scale": ['continuous', [1.0, 1.4]], # only for back type 1
+            "back_x_offset": ['continuous', [-0.1, 0.15]], # back top x length
+            "back_y_offset": ['continuous', [0.0, 0.4]], # back top y coord
+            "back_profile_partial": ['continuous', [0.4, 0.8]], # only for back type 1
+            "back_profile_horizontal_ncuts": ['discrete', [2, 3, 4]], # only for back type 2
+            "back_profile_horizontal_locs0": ['continuous', [1, 2]], # only for back type 2
+            "back_profile_horizontal_locs1": ['continuous', [1, 2]], # only for back type 2
+            "back_profile_horizontal_locs2": ['continuous', [1, 2]], # only for back type 2
+            "back_profile_horizontal_locs3": ['continuous', [1, 2]], # only for back type 2
+            "back_profile_horizontal_ratio": ['continuous', [0.2, 0.8]], # only for back type 2
+            "back_profile_horizontal_lowest": ['continuous', [0, 0.4]], # only for back type 2
+            "back_profile_vertical": ['continuous', [0.8, 0.9]], # only for back type 3
+        }
+    
+    def fix_unused_params(self, params):
+        # check unused parameters inside a given parameter set, and fix them into mid value - for training
+        if params['leg_type'] != 2 and params['leg_type'] != 3:
+            params['limb_profile'] = (self.params_dict['limb_profile'][1][0] + self.params_dict['limb_profile'][1][-1]) / 2
+        if params['has_leg_x_bar'] == 0:
+            params['leg_offset_bar1'] = (self.params_dict['leg_offset_bar1'][1][0] + self.params_dict['leg_offset_bar1'][1][-1]) / 2
+        if params['has_leg_y_bar'] == 0:
+            params['leg_offset_bar0'] = (self.params_dict['leg_offset_bar0'][1][0] + self.params_dict['leg_offset_bar0'][1][-1]) / 2
+        if params['has_arm'] == 0:
+            params['arm_thickness'] = (self.params_dict['arm_thickness'][1][0] + self.params_dict['arm_thickness'][1][-1]) / 2
+            params['arm_height'] = (self.params_dict['arm_height'][1][0] + self.params_dict['arm_height'][1][-1]) / 2
+            params['arm_y'] = (self.params_dict['arm_y'][1][0] + self.params_dict['arm_y'][1][-1]) / 2
+            params['arm_z'] = (self.params_dict['arm_z'][1][0] + self.params_dict['arm_z'][1][-1]) / 2
+            params['arm_mid0'] = (self.params_dict['arm_mid0'][1][0] + self.params_dict['arm_mid0'][1][-1]) / 2
+            params['arm_mid1'] = (self.params_dict['arm_mid1'][1][0] + self.params_dict['arm_mid1'][1][-1]) / 2
+            params['arm_mid2'] = (self.params_dict['arm_mid2'][1][0] + self.params_dict['arm_mid2'][1][-1]) / 2
+            params['arm_profile0'] = (self.params_dict['arm_profile0'][1][0] + self.params_dict['arm_profile0'][1][-1]) / 2
+            params['arm_profile1'] = (self.params_dict['arm_profile1'][1][0] + self.params_dict['arm_profile1'][1][-1]) / 2
+        if params['back_type'] != 3:
+            params['back_vertical_cuts'] = (self.params_dict['back_vertical_cuts'][1][0] + self.params_dict['back_vertical_cuts'][1][-1]) / 2
+            params['back_profile_vertical'] = (self.params_dict['back_profile_vertical'][1][0] + self.params_dict['back_profile_vertical'][1][-1]) / 2
+        if params['back_type'] != 2:
+            params['back_profile_horizontal_ncuts'] = (self.params_dict['back_profile_horizontal_ncuts'][1][0] + self.params_dict['back_profile_horizontal_ncuts'][1][-1]) / 2
+            params['back_profile_horizontal_locs0'] = (self.params_dict['back_profile_horizontal_locs0'][1][0] + self.params_dict['back_profile_horizontal_locs0'][1][-1]) / 2
+            params['back_profile_horizontal_locs1'] = (self.params_dict['back_profile_horizontal_locs1'][1][0] + self.params_dict['back_profile_horizontal_locs1'][1][-1]) / 2
+            params['back_profile_horizontal_locs2'] = (self.params_dict['back_profile_horizontal_locs2'][1][0] + self.params_dict['back_profile_horizontal_locs2'][1][-1]) / 2
+            params['back_profile_horizontal_ratio'] = (self.params_dict['back_profile_horizontal_ratio'][1][0] + self.params_dict['back_profile_horizontal_ratio'][1][-1]) / 2
+            params['back_profile_horizontal_lowest'] = (self.params_dict['back_profile_horizontal_lowest'][1][0] + self.params_dict['back_profile_horizontal_lowest'][1][-1]) / 2
+        if params['back_type'] != 1:
+            params['back_partial_scale'] = (self.params_dict['back_partial_scale'][1][0] + self.params_dict['back_partial_scale'][1][-1]) / 2
+            params['back_profile_partial'] = (self.params_dict['back_profile_partial'][1][0] + self.params_dict['back_profile_partial'][1][-1]) / 2
+        return params
+    
+    def update_params(self, new_params):
+        # replace the parameters and calculate all the new values
+        self.width = new_params["width"]
+        self.size = new_params["size"]
+        self.thickness = new_params["thickness"]
+        self.bevel_width = self.thickness * new_params["bevel_width"]
+        self.seat_back = new_params["seat_back"]
+        self.seat_mid = new_params["seat_mid"]
+        self.seat_mid_x = uniform(
+            self.seat_back + self.seat_mid * (1 - self.seat_back), 1
+        )
+        self.seat_mid_z = new_params["seat_mid_z"]
+        self.seat_front = new_params["seat_front"]
+        self.is_seat_round = new_params["is_seat_round"]
+        self.is_seat_subsurf = new_params["is_seat_subsurf"]
+
+        self.leg_thickness = new_params["leg_thickness"]
+        self.limb_profile = new_params["limb_profile"]
+        self.leg_height = new_params["leg_height"]
+        self.back_height = new_params["back_height"]
+        self.is_leg_round = new_params["is_leg_round"]
+        self.leg_type = self.leg_types[new_params["leg_type"]]
+
+        self.leg_x_offset = 0
+        self.leg_y_offset = 0, 0
+        self.back_x_offset = 0
+        self.back_y_offset = 0
+
+        self.has_leg_x_bar = new_params["has_leg_x_bar"]
+        self.has_leg_y_bar = new_params["has_leg_y_bar"]
+        self.leg_offset_bar = new_params["leg_offset_bar0"], new_params["leg_offset_bar1"]
+
+        self.has_arm = new_params["has_arm"]
+        self.arm_thickness = new_params["arm_thickness"]
+        self.arm_height = self.arm_thickness * new_params["arm_height"]
+        self.arm_y = new_params["arm_y"] * self.size
+        self.arm_z = new_params["arm_z"] * self.back_height
+        self.arm_mid = np.array(
+            [new_params["arm_mid0"], new_params["arm_mid1"], new_params["arm_mid2"]]
+        )
+        self.arm_profile = (new_params["arm_profile0"], new_params["arm_profile1"])
+
+        self.back_thickness = new_params["back_thickness"]
+        self.back_type = self.back_types[new_params["back_type"]]
+        self.back_profile = [(0, 1)]
+        self.back_vertical_cuts = new_params["back_vertical_cuts"]
+        self.back_partial_scale = new_params["back_partial_scale"]
+
+        if self.leg_type == "vertical":
+            self.leg_x_offset = 0
+            self.leg_y_offset = 0, 0
+            self.back_x_offset = 0
+            self.back_y_offset = 0
+        else:
+            self.leg_x_offset = self.width * new_params["leg_x_offset"]
+            self.leg_y_offset = self.size * np.array([new_params["leg_y_offset0"], new_params["leg_y_offset1"]])
+            self.back_x_offset = self.width * new_params["back_x_offset"]
+            self.back_y_offset = self.size * new_params["back_y_offset"]
+
+        match self.back_type:
+            case "partial":
+                self.back_profile = ((new_params["back_profile_partial"], 1),)
+            case "horizontal-bar":
+                n_cuts = int(new_params["back_profile_horizontal_ncuts"])
+                locs = np.array([new_params["back_profile_horizontal_locs0"], new_params["back_profile_horizontal_locs1"],
+                                    new_params["back_profile_horizontal_locs2"], new_params["back_profile_horizontal_locs3"]])[:n_cuts].cumsum()
+                locs = locs / locs[-1]
+                ratio = new_params["back_profile_horizontal_ratio"]
+                locs = np.array(
+                    [
+                        (p + ratio * (l - p), l)
+                        for p, l in zip([0, *locs[:-1]], locs)
+                    ]
+                )
+                lowest = new_params["back_profile_horizontal_lowest"]
+                self.back_profile = locs * (1 - lowest) + lowest
+            case "vertical-bar":
+                self.back_profile = ((new_params["back_profile_vertical"], 1),)
+            case _:
+                self.back_profile = [(0, 1)]
+
+        # TODO: handle the material into the optimization loop
+        materials = AssetList["ChairFactory"]()
+        self.limb_surface = materials["limb"].assign_material()
+        self.surface = materials["surface"].assign_material()
+        if uniform() < 0.3:
+            self.panel_surface = self.surface
+        else:
+            self.panel_surface = materials["panel"].assign_material()
+
+        scratch_prob, edge_wear_prob = materials["wear_tear_prob"]
+        self.scratch, self.edge_wear = materials["wear_tear"]
+        is_scratch = uniform() < scratch_prob
+        is_edge_wear = uniform() < edge_wear_prob
+        if not is_scratch:
+            self.scratch = None
+        if not is_edge_wear:
+            self.edge_wear = None
+
+        # from infinigen.assets.clothes import blanket
+        # from infinigen.assets.scatters.clothes import ClothesCover
+        # self.clothes_scatter = ClothesCover(factory_fn=blanket.BlanketFactory, width=log_uniform(.8, 1.2),
+        #                                    size=uniform(.8, 1.2)) if uniform() < .3 else NoApply()
+        self.clothes_scatter = NoApply()
+        
+
+    def post_init(self):
+        with FixedSeed(self.factory_seed):
+            if self.leg_type == "vertical":
+                self.leg_x_offset = 0
+                self.leg_y_offset = 0, 0
+                self.back_x_offset = 0
+                self.back_y_offset = 0
+            else:
+                self.leg_x_offset = self.width * uniform(0.05, 0.2)
+                self.leg_y_offset = self.size * uniform(0.05, 0.2, 2)
+                self.back_x_offset = self.width * uniform(-0.1, 0.15)
+                self.back_y_offset = self.size * uniform(0.1, 0.25)
+
+            match self.back_type:
+                case "partial":
+                    self.back_profile = ((uniform(0.4, 0.8), 1),)
+                case "horizontal-bar":
+                    n_cuts = np.random.randint(2, 4)
+                    locs = uniform(1, 2, n_cuts).cumsum()
+                    locs = locs / locs[-1]
+                    ratio = uniform(0.5, 0.75)
+                    locs = np.array(
+                        [
+                            (p + ratio * (l - p), l)
+                            for p, l in zip([0, *locs[:-1]], locs)
+                        ]
+                    )
+                    lowest = uniform(0, 0.4)
+                    self.back_profile = locs * (1 - lowest) + lowest
+                case "vertical-bar":
+                    self.back_profile = ((uniform(0.8, 0.9), 1),)
+                case _:
+                    self.back_profile = [(0, 1)]
+
+    def create_placeholder(self, **kwargs) -> bpy.types.Object:
+        obj = new_bbox(
+            -self.width / 2 - max(self.leg_x_offset, self.back_x_offset),
+            self.width / 2 + max(self.leg_x_offset, self.back_x_offset),
+            -self.size - self.leg_y_offset[1] - self.leg_thickness * 0.5,
+            max(self.leg_y_offset[0], self.back_y_offset),
+            -self.leg_height,
+            self.back_height * 1.2,
+        )
+        obj.rotation_euler.z += np.pi / 2
+        butil.apply_transform(obj)
+        return obj
+
+    def create_asset(self, **params) -> bpy.types.Object:
+        obj = self.make_seat()
+        legs = self.make_legs()
+        backs = self.make_backs()
+
+        parts = [obj] + legs + backs
+        parts.extend(self.make_leg_decors(legs))
+        if self.has_arm:
+            parts.extend(self.make_arms(obj, backs))
+        parts.extend(self.make_back_decors(backs))
+
+        for obj in legs:
+            self.solidify(obj, 2)
+        for obj in backs:
+            self.solidify(obj, 2, self.back_thickness)
+
+        obj = join_objects(parts)
+        obj.rotation_euler.z += np.pi / 2
+        butil.apply_transform(obj)
+
+        with FixedSeed(self.factory_seed):
+            # TODO: wasteful to create unique materials for each individual asset
+            self.surface.apply(obj)
+            self.panel_surface.apply(obj, selection="panel")
+            self.limb_surface.apply(obj, selection="limb")
+
+        return obj
+
+    def finalize_assets(self, assets):
+        if self.scratch:
+            self.scratch.apply(assets)
+        if self.edge_wear:
+            self.edge_wear.apply(assets)
+
+    def make_seat(self):
+        x_anchors = (
+            np.array(
+                [
+                    0,
+                    -self.seat_back,
+                    -self.seat_mid_x,
+                    -1,
+                    0,
+                    1,
+                    self.seat_mid_x,
+                    self.seat_back,
+                    0,
+                ]
+            )
+            * self.width
+            / 2
+        )
+        y_anchors = (
+            np.array(
+                [0, 0, -self.seat_mid, -1, -self.seat_front, -1, -self.seat_mid, 0, 0]
+            )
+            * self.size
+        )
+        z_anchors = (
+            np.array([0, 0, self.seat_mid_z, 0, 0, 0, self.seat_mid_z, 0, 0])
+            * self.thickness
+        )
+        vector_locations = [1, 7] if self.is_seat_round else [1, 3, 5, 7]
+        obj = bezier_curve((x_anchors, y_anchors, z_anchors), vector_locations, 8)
+        with butil.ViewportMode(obj, "EDIT"):
+            bpy.ops.mesh.select_all(action="SELECT")
+            bpy.ops.mesh.fill_grid(use_interp_simple=True)
+        butil.modify_mesh(obj, "SOLIDIFY", thickness=self.thickness, offset=0)
+        subsurf(obj, 1, not self.is_seat_subsurf)
+        butil.modify_mesh(obj, "BEVEL", width=self.bevel_width, segments=8)
+        return obj
+
+    def make_legs(self):
+        leg_starts = np.array(
+            [[-self.seat_back, 0, 0], [-1, -1, 0], [1, -1, 0], [self.seat_back, 0, 0]]
+        ) * np.array([[self.width / 2, self.size, 0]])
+        leg_ends = leg_starts.copy()
+        leg_ends[[0, 1], 0] -= self.leg_x_offset
+        leg_ends[[2, 3], 0] += self.leg_x_offset
+        leg_ends[[0, 3], 1] += self.leg_y_offset[0]
+        leg_ends[[1, 2], 1] -= self.leg_y_offset[1]
+        leg_ends[:, -1] = -self.leg_height
+        return self.make_limb(leg_ends, leg_starts)
+
+    def make_limb(self, leg_ends, leg_starts):
+        limbs = []
+        for leg_start, leg_end in zip(leg_starts, leg_ends):
+            match self.leg_type:
+                case "up-curved":
+                    axes = [(0, 0, 1), None]
+                    scale = [self.limb_profile, 1]
+                case "down-curved":
+                    axes = [None, (0, 0, 1)]
+                    scale = [1, self.limb_profile]
+                case _:
+                    axes = None
+                    scale = None
+            limb = align_bezier(
+                np.stack([leg_start, leg_end], -1), axes, scale, resolution=64
+            )
+            limb.location = (
+                np.array(
+                    [
+                        1 if leg_start[0] < 0 else -1,
+                        1 if leg_start[1] < -self.size / 2 else -1,
+                        0,
+                    ]
+                )
+                * self.leg_thickness
+                / 2
+            )
+            butil.apply_transform(limb, True)
+            limbs.append(limb)
+        return limbs
+
+    def make_backs(self):
+        back_starts = (
+            np.array([[-self.seat_back, 0, 0], [self.seat_back, 0, 0]]) * self.width / 2
+        )
+        back_ends = back_starts.copy()
+        back_ends[:, 0] += np.array([self.back_x_offset, -self.back_x_offset])
+        back_ends[:, 1] = self.back_y_offset
+        back_ends[:, 2] = self.back_height
+        return self.make_limb(back_starts, back_ends)
+
+    def make_leg_decors(self, legs):
+        decors = []
+        if self.has_leg_x_bar:
+            z_height = -self.leg_height * uniform(*self.leg_offset_bar)
+            locs = []
+            for leg in legs:
+                co = read_co(leg)
+                locs.append(co[np.argmin(np.abs(co[:, -1] - z_height))])
+            decors.append(
+                self.solidify(bezier_curve(np.stack([locs[0], locs[3]], -1)), 0)
+            )
+            decors.append(
+                self.solidify(bezier_curve(np.stack([locs[1], locs[2]], -1)), 0)
+            )
+        if self.has_leg_y_bar:
+            z_height = -self.leg_height * uniform(*self.leg_offset_bar)
+            locs = []
+            for leg in legs:
+                co = read_co(leg)
+                locs.append(co[np.argmin(np.abs(co[:, -1] - z_height))])
+            decors.append(
+                self.solidify(bezier_curve(np.stack([locs[0], locs[1]], -1)), 1)
+            )
+            decors.append(
+                self.solidify(bezier_curve(np.stack([locs[2], locs[3]], -1)), 1)
+            )
+        for d in decors:
+            write_attribute(d, 1, "limb", "FACE")
+        return decors
+
+    def make_back_decors(self, backs, finalize=True):
+        obj = join_objects([deep_clone_obj(b) for b in backs])
+        x, y, z = read_co(obj).T
+        x += np.where(x > 0, self.back_thickness / 2, -self.back_thickness / 2)
+        write_co(obj, np.stack([x, y, z], -1))
+        smoothness = uniform(0, 1)
+        profile_shape_factor = uniform(0, 0.4)
+        with butil.ViewportMode(obj, "EDIT"):
+            bpy.ops.mesh.select_mode(type="EDGE")
+            center = read_edge_center(obj)
+            for z_min, z_max in self.back_profile:
+                select_edges(
+                    obj,
+                    (z_min * self.back_height <= center[:, -1])
+                    & (center[:, -1] <= z_max * self.back_height),
+                )
+                bpy.ops.mesh.bridge_edge_loops(
+                    number_cuts=32,
+                    interpolation="LINEAR",
+                    smoothness=smoothness,
+                    profile_shape_factor=profile_shape_factor,
+                )
+            bpy.ops.mesh.select_loose()
+            bpy.ops.mesh.delete()
+        butil.modify_mesh(
+            obj,
+            "SOLIDIFY",
+            thickness=np.minimum(self.thickness, self.back_thickness),
+            offset=0,
+        )
+        if finalize:
+            butil.modify_mesh(obj, "BEVEL", width=self.bevel_width, segments=8)
+        parts = [obj]
+        if self.back_type == "vertical-bar":
+            other = join_objects([deep_clone_obj(b) for b in backs])
+            with butil.ViewportMode(other, "EDIT"):
+                bpy.ops.mesh.select_mode(type="EDGE")
+                bpy.ops.mesh.select_all(action="SELECT")
+                bpy.ops.mesh.bridge_edge_loops(
+                    number_cuts=self.back_vertical_cuts,
+                    interpolation="LINEAR",
+                    smoothness=smoothness,
+                    profile_shape_factor=profile_shape_factor,
+                )
+                bpy.ops.mesh.select_all(action="INVERT")
+                bpy.ops.mesh.delete()
+                bpy.ops.mesh.select_all(action="SELECT")
+                bpy.ops.mesh.delete(type="ONLY_FACE")
+            remove_edges(other, np.abs(read_edge_direction(other)[:, -1]) < 0.5)
+            remove_vertices(other, lambda x, y, z: z < -self.thickness / 2)
+            remove_vertices(
+                other,
+                lambda x, y, z: z
+                > (self.back_profile[0][0] + self.back_profile[0][1])
+                * self.back_height
+                / 2,
+            )
+            parts.append(self.solidify(other, 2, self.back_thickness))
+        elif self.back_type == "partial":
+            co = read_co(obj)
+            co[:, 1] *= self.back_partial_scale
+            write_co(obj, co)
+        for p in parts:
+            write_attribute(p, 1, "panel", "FACE")
+        return parts
+
+    def make_arms(self, base, backs):
+        co = read_co(base)
+        end = co[np.argmin(co[:, 0] - (np.abs(co[:, 1] + self.arm_y) < 0.02))]
+        end[0] += self.arm_thickness / 4
+        end_ = end.copy()
+        end_[0] = -end[0]
+        arms = []
+        co = read_co(backs[0])
+        start = co[np.argmin(co[:, 0] - (np.abs(co[:, -1] - self.arm_z) < 0.02))]
+        start[0] -= self.arm_thickness / 4
+        start_ = start.copy()
+        start_[0] = -start[0]
+        for start, end in zip([start, start_], [end, end_]):
+            mid = np.array(
+                [
+                    end[0] + self.arm_mid[0] * (-1 if end[0] > 0 else 1),
+                    end[1] + self.arm_mid[1],
+                    start[2] + self.arm_mid[2],
+                ]
+            )
+            arm = align_bezier(
+                np.stack([start, mid, end], -1),
+                np.array(
+                    [
+                        [end[0] - start[0], end[1] - start[1], 0],
+                        [0, 1 / np.sqrt(2), 1 / np.sqrt(2)],
+                        [0, 0, 1],
+                    ]
+                ),
+                [1, *self.arm_profile, 1],
+            )
+            if self.is_leg_round:
+                surface.add_geomod(
+                    arm,
+                    geo_radius,
+                    apply=True,
+                    input_args=[self.arm_thickness / 2, 32],
+                    input_kwargs={"to_align_tilt": False},
+                )
+            else:
+                with butil.ViewportMode(arm, "EDIT"):
+                    bpy.ops.mesh.select_all(action="SELECT")
+                    bpy.ops.mesh.extrude_edges_move(
+                        TRANSFORM_OT_translate={
+                            "value": (
+                                self.arm_thickness
+                                if end[0] < 0
+                                else -self.arm_thickness,
+                                0,
+                                0,
+                            )
+                        }
+                    )
+                butil.modify_mesh(arm, "SOLIDIFY", thickness=self.arm_height, offset=0)
+            write_attribute(arm, 1, "limb", "FACE")
+            arms.append(arm)
+        return arms
+
+    def solidify(self, obj, axis, thickness=None):
+        if thickness is None:
+            thickness = self.leg_thickness
+        if self.is_leg_round:
+            solidify(obj, axis, thickness)
+            butil.modify_mesh(obj, "BEVEL", width=self.bevel_width, segments=8)
+        else:
+            surface.add_geomod(
+                obj, geo_radius, apply=True, input_args=[thickness / 2, 32]
+            )
+        write_attribute(obj, 1, "limb", "FACE")
+        return obj

core/assets/dandelion.py

@@ -0,0 +1,1097 @@
+# Copyright (C) 2023, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Beining Han
+# Acknowledgement: This file draws inspiration from https://www.youtube.com/watch?v=61Sk8j1Ml9c by BradleyAnimation
+
+import bpy
+import numpy as np
+from numpy.random import normal, randint, uniform
+
+import infinigen
+from infinigen.assets.materials import simple_brownish, simple_greenery, simple_whitish
+from infinigen.core import surface
+from infinigen.core.nodes import node_utils
+from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.tagging import tag_nodegroup, tag_object
+from infinigen.core.util.math import FixedSeed
+
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_head_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_head_geometry(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketVectorTranslation", "Translation", (0.0, 0.0, 1.0)),
+            ("NodeSocketFloatDistance", "Radius", 0.04),
+        ],
+    )
+
+    uv_sphere_1 = nw.new_node(
+        Nodes.MeshUVSphere,
+        input_kwargs={"Segments": 64, "Radius": group_input.outputs["Radius"]},
+    )
+
+    transform_1 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": uv_sphere_1,
+            "Translation": group_input.outputs["Translation"],
+        },
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": transform_1,
+            "Material": surface.shaderfunc_to_material(
+                simple_brownish.shader_simple_brown
+            ),
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": set_material}
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_end_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_end_geometry(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketGeometry", "Points", None)]
+    )
+
+    endpoint_selection = nw.new_node(
+        "GeometryNodeCurveEndpointSelection", input_kwargs={"End Size": 0}
+    )
+
+    uv_sphere = nw.new_node(
+        Nodes.MeshUVSphere, input_kwargs={"Segments": 64, "Radius": 0.04}
+    )
+
+    vector = nw.new_node(Nodes.Vector)
+    vector.vector = (uniform(0.45, 0.7), uniform(0.45, 0.7), uniform(2, 3))
+
+    transform = nw.new_node(
+        Nodes.Transform, input_kwargs={"Geometry": uv_sphere, "Scale": vector}
+    )
+
+    cone = nw.new_node(
+        "GeometryNodeMeshCone", input_kwargs={"Radius Bottom": 0.0040, "Depth": 0.0040}
+    )
+
+    normal = nw.new_node(Nodes.InputNormal)
+
+    align_euler_to_vector_1 = nw.new_node(
+        Nodes.AlignEulerToVector, input_kwargs={"Vector": normal}, attrs={"axis": "Z"}
+    )
+
+    instance_on_points_1 = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": transform,
+            "Instance": cone.outputs["Mesh"],
+            "Rotation": align_euler_to_vector_1,
+        },
+    )
+
+    join_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [instance_on_points_1, transform]}
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": join_geometry,
+            "Material": surface.shaderfunc_to_material(
+                simple_brownish.shader_simple_brown
+            ),
+        },
+    )
+
+    geometry_to_instance = nw.new_node(
+        "GeometryNodeGeometryToInstance", input_kwargs={"Geometry": set_material}
+    )
+
+    curve_tangent = nw.new_node(Nodes.CurveTangent)
+
+    align_euler_to_vector = nw.new_node(
+        Nodes.AlignEulerToVector,
+        input_kwargs={"Vector": curve_tangent},
+        attrs={"axis": "Z"},
+    )
+
+    instance_on_points = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": group_input.outputs["Points"],
+            "Selection": endpoint_selection,
+            "Instance": geometry_to_instance,
+            "Rotation": align_euler_to_vector,
+        },
+    )
+
+    realize_instances = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": instance_on_points}
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": realize_instances}
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_branch_shape", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_branch_shape(nw: NodeWrangler):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    pedal_stem_branches_num = nw.new_node(
+        Nodes.Integer, label="pedal_stem_branches_num"
+    )
+    pedal_stem_branches_num.integer = 40
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketFloatDistance", "Radius", 0.0100)]
+    )
+
+    curve_circle_1 = nw.new_node(
+        Nodes.CurveCircle,
+        input_kwargs={
+            "Resolution": pedal_stem_branches_num,
+            "Radius": group_input.outputs["Radius"],
+        },
+    )
+
+    pedal_stem_branch_length = nw.new_node(
+        Nodes.Value, label="pedal_stem_branch_length"
+    )
+    pedal_stem_branch_length.outputs[0].default_value = 0.5000
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": pedal_stem_branch_length}
+    )
+
+    curve_line_1 = nw.new_node(Nodes.CurveLine, input_kwargs={"End": combine_xyz_1})
+
+    resample_curve = nw.new_node(
+        Nodes.ResampleCurve, input_kwargs={"Curve": curve_line_1, "Count": 40}
+    )
+
+    spline_parameter = nw.new_node(Nodes.SplineParameter)
+
+    float_curve = nw.new_node(
+        Nodes.FloatCurve, input_kwargs={"Value": spline_parameter.outputs["Factor"]}
+    )
+    node_utils.assign_curve(
+        float_curve.mapping.curves[0],
+        [
+            (0.0000, 0.0000),
+            (0.2, 0.08 * np.random.normal(1.0, 0.15)),
+            (0.4, 0.22 * np.random.normal(1.0, 0.2)),
+            (0.6, 0.45 * np.random.normal(1.0, 0.2)),
+            (0.8, 0.7 * np.random.normal(1.0, 0.1)),
+            (1.0000, 1.0000),
+        ],
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: float_curve, 1: uniform(0.15, 0.4)},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply})
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={"Geometry": resample_curve, "Offset": combine_xyz},
+    )
+
+    normal = nw.new_node(Nodes.InputNormal)
+
+    align_euler_to_vector = nw.new_node(
+        Nodes.AlignEulerToVector, input_kwargs={"Vector": normal}
+    )
+
+    instance_on_points = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": curve_circle_1.outputs["Curve"],
+            "Instance": set_position,
+            "Rotation": align_euler_to_vector,
+        },
+    )
+
+    random_value_1 = nw.new_node(
+        Nodes.RandomValue, input_kwargs={2: -0.2000, 3: 0.2000, "Seed": 2}
+    )
+
+    random_value_2 = nw.new_node(
+        Nodes.RandomValue, input_kwargs={2: -0.2000, 3: 0.2000, "Seed": 1}
+    )
+
+    random_value = nw.new_node(Nodes.RandomValue, input_kwargs={2: -0.2000, 3: 0.2000})
+
+    combine_xyz_2 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={
+            "X": random_value_1.outputs[1],
+            "Y": random_value_2.outputs[1],
+            "Z": random_value.outputs[1],
+        },
+    )
+
+    rotate_instances = nw.new_node(
+        Nodes.RotateInstances,
+        input_kwargs={"Instances": instance_on_points, "Rotation": combine_xyz_2},
+    )
+
+    random_value_3 = nw.new_node(Nodes.RandomValue, input_kwargs={2: 0.8000})
+
+    scale_instances = nw.new_node(
+        Nodes.ScaleInstances,
+        input_kwargs={
+            "Instances": rotate_instances,
+            "Scale": random_value_3.outputs[1],
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Instances": scale_instances},
+        attrs={"is_active_output": True},
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_branch_contour", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_branch_contour(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketGeometry", "Geometry", None)]
+    )
+
+    realize_instances = nw.new_node(
+        Nodes.RealizeInstances,
+        input_kwargs={"Geometry": group_input.outputs["Geometry"]},
+    )
+
+    pedal_stem_branch_rsample = nw.new_node(
+        Nodes.Value, label="pedal_stem_branch_rsample"
+    )
+    pedal_stem_branch_rsample.outputs[0].default_value = 10.0
+
+    resample_curve = nw.new_node(
+        Nodes.ResampleCurve,
+        input_kwargs={"Curve": realize_instances, "Count": pedal_stem_branch_rsample},
+    )
+
+    index = nw.new_node(Nodes.Index)
+
+    capture_attribute = nw.new_node(
+        Nodes.CaptureAttribute,
+        input_kwargs={"Geometry": resample_curve, 5: index},
+        attrs={"domain": "CURVE", "data_type": "INT"},
+    )
+
+    spline_parameter = nw.new_node(Nodes.SplineParameter)
+
+    float_curve = nw.new_node(
+        Nodes.FloatCurve, input_kwargs={"Value": spline_parameter.outputs["Factor"]}
+    )
+
+    # generate pedal branch contour
+    dist = uniform(-0.05, -0.25)
+    node_utils.assign_curve(
+        float_curve.mapping.curves[0],
+        [
+            (0.0, 0.0),
+            (0.2, 0.2 + (dist + normal(0, 0.05)) / 2.0),
+            (0.4, 0.4 + (dist + normal(0, 0.05))),
+            (0.6, 0.6 + (dist + normal(0, 0.05)) / 1.2),
+            (0.8, 0.8 + (dist + normal(0, 0.05)) / 2.4),
+            (1.0, 0.95 + normal(0, 0.05)),
+        ],
+    )
+
+    random_value = nw.new_node(
+        Nodes.RandomValue,
+        input_kwargs={2: 0.05, 3: 0.35, "ID": capture_attribute.outputs[5]},
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: float_curve, 1: random_value.outputs[1]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply})
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={
+            "Geometry": capture_attribute.outputs["Geometry"],
+            "Offset": combine_xyz,
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": set_position}
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_branch_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_branch_geometry(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Curve", None),
+            ("NodeSocketVectorTranslation", "Translation", (0.0, 0.0, 1.0)),
+        ],
+    )
+
+    set_curve_radius_1 = nw.new_node(
+        Nodes.SetCurveRadius,
+        input_kwargs={"Curve": group_input.outputs["Curve"], "Radius": 1.0},
+    )
+
+    curve_circle_2 = nw.new_node(
+        Nodes.CurveCircle,
+        input_kwargs={"Radius": uniform(0.001, 0.0025), "Resolution": 4},
+    )
+
+    curve_to_mesh_1 = nw.new_node(
+        Nodes.CurveToMesh,
+        input_kwargs={
+            "Curve": set_curve_radius_1,
+            "Profile Curve": curve_circle_2.outputs["Curve"],
+            "Fill Caps": True,
+        },
+    )
+
+    transform_2 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": curve_to_mesh_1,
+            "Translation": group_input.outputs["Translation"],
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": transform_2}
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem_geometry(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketVectorTranslation", "End", (0.0, 0.0, 1.0)),
+            ("NodeSocketVectorTranslation", "Middle", (0.0, 0.0, 0.5)),
+            ("NodeSocketFloatDistance", "Radius", 0.05),
+        ],
+    )
+
+    quadratic_bezier = nw.new_node(
+        Nodes.QuadraticBezier,
+        input_kwargs={
+            "Start": (0.0, 0.0, 0.0),
+            "Middle": group_input.outputs["Middle"],
+            "End": group_input.outputs["End"],
+        },
+    )
+
+    set_curve_radius = nw.new_node(
+        Nodes.SetCurveRadius,
+        input_kwargs={
+            "Curve": quadratic_bezier,
+            "Radius": group_input.outputs["Radius"],
+        },
+    )
+
+    curve_circle = nw.new_node(
+        Nodes.CurveCircle, input_kwargs={"Radius": 0.2, "Resolution": 8}
+    )
+
+    curve_to_mesh = nw.new_node(
+        Nodes.CurveToMesh,
+        input_kwargs={
+            "Curve": set_curve_radius,
+            "Profile Curve": curve_circle.outputs["Curve"],
+            "Fill Caps": True,
+        },
+    )
+
+    set_material_2 = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": curve_to_mesh,
+            "Material": surface.shaderfunc_to_material(
+                simple_whitish.shader_simple_white
+            ),
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": set_material_2, "Curve": quadratic_bezier},
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_selection", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_selection(nw: NodeWrangler, params):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    random_value = nw.new_node(Nodes.RandomValue, input_kwargs={5: 1})
+
+    greater_than = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: params["random_dropout"], 1: random_value.outputs[1]},
+        attrs={"operation": "GREATER_THAN"},
+    )
+
+    index_1 = nw.new_node(Nodes.Index)
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketFloat", "num_segments", 0.5)]
+    )
+
+    divide = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: index_1, 1: group_input.outputs["num_segments"]},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    less_than = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: divide, 1: params["row_less_than"]},
+        attrs={"operation": "LESS_THAN"},
+    )
+
+    greater_than_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: divide, 1: params["row_great_than"]},
+        attrs={"operation": "GREATER_THAN"},
+    )
+
+    op_and = nw.new_node(
+        Nodes.BooleanMath, input_kwargs={0: less_than, 1: greater_than_1}
+    )
+
+    modulo = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: index_1, 1: group_input.outputs["num_segments"]},
+        attrs={"operation": "MODULO"},
+    )
+
+    less_than_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: modulo, 1: params["col_less_than"]},
+        attrs={"operation": "LESS_THAN"},
+    )
+
+    greater_than_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: modulo, 1: params["col_great_than"]},
+        attrs={"operation": "GREATER_THAN"},
+    )
+
+    op_and_1 = nw.new_node(
+        Nodes.BooleanMath, input_kwargs={0: less_than_1, 1: greater_than_2}
+    )
+
+    nand = nw.new_node(
+        Nodes.BooleanMath,
+        input_kwargs={0: op_and, 1: op_and_1},
+        attrs={"operation": "NAND"},
+    )
+
+    op_and_2 = nw.new_node(Nodes.BooleanMath, input_kwargs={0: greater_than, 1: nand})
+
+    group_output = nw.new_node(Nodes.GroupOutput, input_kwargs={"Boolean": op_and_2})
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_stem_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_stem_geometry(nw: NodeWrangler, params):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, 
+        expose_input=[
+            ("NodeSocketGeometry", "Curve", None),
+        ]
+    )
+
+    spline_parameter = nw.new_node(Nodes.SplineParameter)
+
+    value = nw.new_node(Nodes.Value)
+    value.outputs[0].default_value = params["stem_map_range"]
+
+    map_range = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={"Value": spline_parameter.outputs["Factor"], 3: 0.4, 4: value},
+    )
+
+    set_curve_radius_2 = nw.new_node(
+        Nodes.SetCurveRadius,
+        input_kwargs={
+            "Curve": group_input.outputs["Curve"],
+            "Radius": map_range.outputs["Result"],
+        },
+    )
+
+    stem_radius = nw.new_node(Nodes.Value, label="stem_radius")
+    stem_radius.outputs[0].default_value = params["stem_radius"]
+
+    curve_circle_3 = nw.new_node(
+        Nodes.CurveCircle, input_kwargs={"Radius": stem_radius}
+    )
+
+    curve_to_mesh_2 = nw.new_node(
+        Nodes.CurveToMesh,
+        input_kwargs={
+            "Curve": set_curve_radius_2,
+            "Profile Curve": curve_circle_3.outputs["Curve"],
+            "Fill Caps": True,
+        },
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": curve_to_mesh_2,
+            "Material": surface.shaderfunc_to_material(
+                simple_greenery.shader_simple_greenery
+            ),
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Mesh": tag_nodegroup(nw, set_material, "stem")},
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_pedal_stem", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_pedal_stem(nw: NodeWrangler, params):
+    # Code generated using version 2.4.3 of the node_transpiler
+    pedal_stem_top_point = nw.new_node(Nodes.Vector, label="pedal_stem_top_point")
+    pedal_stem_top_point.vector = (0.0, 0.0, 1.0)
+
+    pedal_stem_mid_point = nw.new_node(Nodes.Vector, label="pedal_stem_mid_point")
+    pedal_stem_mid_point.vector = (
+        params["pedal_stem_mid_point_x"], 
+        params["pedal_stem_mid_point_y"], 
+        0.5
+    )
+
+    pedal_stem_radius = nw.new_node(Nodes.Value, label="pedal_stem_radius")
+    pedal_stem_radius.outputs[0].default_value = params["pedal_stem_radius"]
+
+    pedal_stem_geometry = nw.new_node(
+        nodegroup_pedal_stem_geometry().name,
+        input_kwargs={
+            "End": pedal_stem_top_point,
+            "Middle": pedal_stem_mid_point,
+            "Radius": pedal_stem_radius,
+        },
+    )
+
+    pedal_stem_top_radius = nw.new_node(Nodes.Value, label="pedal_stem_top_radius")
+    pedal_stem_top_radius.outputs[0].default_value = params["pedal_stem_top_radius"]
+
+    pedal_stem_branch_shape = nw.new_node(
+        nodegroup_pedal_stem_branch_shape().name,
+        input_kwargs={"Radius": pedal_stem_top_radius},
+    )
+
+    pedal_stem_branch_geometry = nw.new_node(
+        nodegroup_pedal_stem_branch_geometry().name,
+        input_kwargs={
+            "Curve": pedal_stem_branch_shape,
+            "Translation": pedal_stem_top_point,
+        },
+    )
+
+    set_material_3 = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": pedal_stem_branch_geometry,
+            "Material": surface.shaderfunc_to_material(
+                simple_whitish.shader_simple_white
+            ),
+        },
+    )
+
+    resample_curve = nw.new_node(
+        Nodes.ResampleCurve,
+        input_kwargs={"Curve": pedal_stem_geometry.outputs["Curve"]},
+    )
+
+    pedal_stem_end_geometry = nw.new_node(
+        nodegroup_pedal_stem_end_geometry().name,
+        input_kwargs={"Points": resample_curve},
+    )
+
+    pedal_stem_head_geometry = nw.new_node(
+        nodegroup_pedal_stem_head_geometry().name,
+        input_kwargs={
+            "Translation": pedal_stem_top_point,
+            "Radius": pedal_stem_top_radius,
+        },
+    )
+
+    join_geometry = nw.new_node(
+        Nodes.JoinGeometry,
+        input_kwargs={
+            "Geometry": [
+                pedal_stem_geometry.outputs["Geometry"],
+                set_material_3,
+                pedal_stem_end_geometry,
+                pedal_stem_head_geometry,
+            ]
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": join_geometry}
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_flower_geometry", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_flower_geometry(nw: NodeWrangler, params):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    num_core_segments = nw.new_node(
+        Nodes.Integer, label="num_core_segments", attrs={"integer": 10}
+    )
+    num_core_segments.integer = params["flower_num_core_segments"]
+
+    num_core_rings = nw.new_node(
+        Nodes.Integer, label="num_core_rings", attrs={"integer": 10}
+    )
+    num_core_rings.integer = params["flower_num_core_rings"]
+
+    uv_sphere_2 = nw.new_node(
+        Nodes.MeshUVSphere,
+        input_kwargs={
+            "Segments": num_core_segments,
+            "Rings": num_core_rings,
+            "Radius": params["flower_radius"],
+        },
+    )
+
+    flower_core_shape = nw.new_node(Nodes.Vector, label="flower_core_shape")
+    flower_core_shape.vector = (params["flower_core_shape_x"], params["flower_core_shape_y"], params["flower_core_shape_z"])
+
+    transform = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={"Geometry": uv_sphere_2, "Scale": flower_core_shape},
+    )
+
+    selection_params = {
+        "random_dropout": params["random_dropout"],
+        "row_less_than": int(params["row_less_than"] * num_core_rings.integer),
+        "row_great_than": int(params["row_great_than"] * num_core_rings.integer),
+        "col_less_than": int(params["col_less_than"] * num_core_segments.integer),
+        "col_great_than": int(params["col_less_than"] * num_core_segments.integer),
+    }
+    pedal_selection = nw.new_node(
+        nodegroup_pedal_selection(params=selection_params).name,
+        input_kwargs={"num_segments": num_core_segments},
+    )
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketGeometry", "Instance", None)]
+    )
+
+    normal_1 = nw.new_node(Nodes.InputNormal)
+
+    align_euler_to_vector_1 = nw.new_node(
+        Nodes.AlignEulerToVector, input_kwargs={"Vector": normal_1}, attrs={"axis": "Z"}
+    )
+
+    random_value_1 = nw.new_node(Nodes.RandomValue, input_kwargs={2: 0.4, 3: 0.7})
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: random_value_1.outputs[1]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    instance_on_points_1 = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": transform,
+            "Selection": pedal_selection,
+            "Instance": group_input.outputs["Instance"],
+            "Rotation": align_euler_to_vector_1,
+            "Scale": multiply,
+        },
+    )
+
+    realize_instances_1 = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": instance_on_points_1}
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={
+            "Geometry": transform,
+            "Material": surface.shaderfunc_to_material(
+                simple_whitish.shader_simple_white
+            ),
+        },
+    )
+
+    join_geometry_1 = nw.new_node(
+        Nodes.JoinGeometry,
+        input_kwargs={"Geometry": [realize_instances_1, set_material]},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": tag_nodegroup(nw, join_geometry_1, "flower")},
+    )
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_flower_on_stem", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_flower_on_stem(nw: NodeWrangler):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Points", None),
+            ("NodeSocketGeometry", "Instance", None),
+        ],
+    )
+
+    endpoint_selection = nw.new_node(
+        "GeometryNodeCurveEndpointSelection", input_kwargs={"Start Size": 0}
+    )
+
+    curve_tangent = nw.new_node(Nodes.CurveTangent)
+
+    align_euler_to_vector_2 = nw.new_node(
+        Nodes.AlignEulerToVector,
+        input_kwargs={"Vector": curve_tangent},
+        attrs={"axis": "Z"},
+    )
+
+    instance_on_points_2 = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": group_input.outputs["Points"],
+            "Selection": endpoint_selection,
+            "Instance": group_input.outputs["Instance"],
+            "Rotation": align_euler_to_vector_2,
+        },
+    )
+
+    realize_instances_2 = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": instance_on_points_2}
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Instances": realize_instances_2}
+    )
+
+
+def geometry_dandelion_nodes(nw: NodeWrangler, **kwargs):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    quadratic_bezier_1 = nw.new_node(
+        Nodes.QuadraticBezier,
+        input_kwargs={
+            "Start": (0.0, 0.0, 0.0),
+            "Middle": (kwargs["bezier_middle_x"], kwargs["bezier_middle_y"], 0.5),
+            "End": (kwargs["bezier_end_x"], kwargs["bezier_end_y"], 1.0),
+        },
+    )
+
+    resample_curve = nw.new_node(
+        Nodes.ResampleCurve, input_kwargs={"Curve": quadratic_bezier_1}
+    )
+
+    pedal_stem = nw.new_node(
+        nodegroup_pedal_stem(kwargs).name, 
+        input_kwargs={},
+    )
+
+    geometry_to_instance = nw.new_node(
+        "GeometryNodeGeometryToInstance", input_kwargs={"Geometry": pedal_stem}
+    )
+
+    flower_geometry = nw.new_node(
+        nodegroup_flower_geometry(kwargs).name,
+        input_kwargs={"Instance": geometry_to_instance},
+    )
+
+    geometry_to_instance_1 = nw.new_node(
+        "GeometryNodeGeometryToInstance", input_kwargs={"Geometry": flower_geometry}
+    )
+
+    value_2 = nw.new_node(Nodes.Value)
+    value_2.outputs[0].default_value = kwargs["transform_scale"]
+
+    transform_3 = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={"Geometry": geometry_to_instance_1, "Scale": value_2},
+    )
+
+    flower_on_stem = nw.new_node(
+        nodegroup_flower_on_stem().name,
+        input_kwargs={"Points": resample_curve, "Instance": transform_3},
+    )
+
+    stem_geometry = nw.new_node(
+        nodegroup_stem_geometry(kwargs).name, 
+        input_kwargs={
+            "Curve": quadratic_bezier_1,
+        }
+    )
+
+    join_geometry_2 = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [flower_on_stem, stem_geometry]}
+    )
+
+    realize_instances = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": join_geometry_2}
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": realize_instances}
+    )
+
+
+def geometry_dandelion_seed_nodes(nw: NodeWrangler, **kwargs):
+    # Code generated using version 2.4.3 of the node_transpiler
+
+    pedal_stem = nw.new_node(nodegroup_pedal_stem().name)
+
+    geometry_to_instance = nw.new_node(
+        "GeometryNodeGeometryToInstance", input_kwargs={"Geometry": pedal_stem}
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": geometry_to_instance}
+    )
+
+flower_modes_dict = {
+    0: "full_flower",
+    1: "no_flower",
+    2: "sparse_flower",
+}
+class DandelionFactory(AssetFactory):
+    def __init__(self, factory_seed, coarse=False):
+        super(DandelionFactory, self).__init__(factory_seed, coarse=coarse)
+        self.get_params_dict()
+
+        with FixedSeed(factory_seed):
+            self.sample_parameters()
+    
+    def get_params_dict(self):
+        # list all the parameters (key:name, value: [type, range]) used in this generator
+        self.params_dict = {
+            "flower_mode": ["discrete", (0, 1, 2)],
+            "random_dropout": ["continuous", (0.2, 0.6)],
+            "row_less_than": ["continuous", (0.0, 1.0)],
+            "col_less_than": ["continuous", (0.0, 1.0)],
+            "row_great_than": ["continuous", (0.0, 1.0)],
+            "col_great_than": ["continuous", (0.0, 1.0)],
+            "bezier_middle_x": ["continuous", (-0.6, 0.6)],
+            "bezier_middle_y": ["continuous", (-0.6, 0.6)],
+            "bezier_end_x": ["continuous", (-0.6, 0.6)],
+            "bezier_end_y": ["continuous", (-0.6, 0.6)],
+            "flower_num_core_segments": ["discrete", (8, 15, 20, 25)],
+            "flower_num_core_rings": ["discrete", (8, 15, 20)],
+            "transform_scale": ["continuous", (-0.7, -0.1)],
+            "stem_map_range": ["continuous", (0.1, 0.6)],
+            "stem_radius": ["continuous", (0.01, 0.03)],
+        }
+    
+    def sample_parameters(self):
+        # sample all the parameters
+        flower_mode = flower_modes_dict[randint(0, 2)]
+        if flower_mode == "full_flower":
+            random_dropout = 1.0
+            row_less_than = 0.0
+            row_great_than = 0.0
+            col_less_than = 0.0
+            col_great_than = 0.0
+        elif flower_mode == "no_flower":
+            random_dropout = 0.0
+            row_less_than = 1.0
+            row_great_than = 0.0
+            col_less_than = 1.0
+            col_great_than = 0.0
+        elif flower_mode == "sparse_flower":
+            random_dropout = uniform(0.2, 0.6)
+            row_less_than = 0.0
+            row_great_than = 0.0
+            col_less_than = 0.0
+            col_great_than = 0.0
+        else:
+            raise ValueError("Invalid flower mode")
+        self.params = {
+            "flower_mode": flower_mode,
+            "random_dropout": random_dropout,
+            "row_less_than": row_less_than,
+            "row_great_than": row_great_than,
+            "col_less_than": col_less_than,
+            "col_great_than": col_great_than,
+            "bezier_middle_x": normal(0.0, 0.1),
+            "bezier_middle_y": normal(0.0, 0.1),
+            "bezier_end_x": normal(0.0, 0.1),
+            "bezier_end_y": normal(0.0, 0.1),
+            "pedal_stem_mid_point_x": normal(0.0, 0.05),
+            "pedal_stem_mid_point_y": normal(0.0, 0.05),
+            "pedal_stem_radius": uniform(0.02, 0.045),
+            "pedal_stem_top_radius": uniform(0.005, 0.008),
+            "flower_num_core_segments": randint(8, 25),
+            "flower_num_core_rings": randint(8, 20),
+            "flower_radius": uniform(0.02, 0.05),
+            "flower_core_shape_x": uniform(0.8, 1.2),
+            "flower_core_shape_y": uniform(0.8, 1.2),
+            "flower_core_shape_z": uniform(0.5, 0.8),
+            "transform_scale": uniform(-0.5, -0.15),
+            "stem_map_range": uniform(0.2, 0.4),
+            "stem_radius": uniform(0.01, 0.024),
+        }
+    
+    def fix_unused_params(self, params):
+        return params
+
+    def update_params(self, params):
+        # update the parameters in the node graph
+        flower_mode = flower_modes_dict[params["flower_mode"]]
+        if flower_mode == "full_flower":
+            random_dropout = uniform(0.7, 1.0)
+            row_less_than = 0.0
+            row_great_than = 0.0
+            col_less_than = 0.0
+            col_great_than = 0.0
+        elif flower_mode == "no_flower":
+            random_dropout = 0.0
+            row_less_than = 1.0
+            row_great_than = 0.0
+            col_less_than = 1.0
+            col_great_than = 0.0
+        elif flower_mode == "sparse_flower":
+            random_dropout = params["random_dropout"]
+            row_less_than = params["row_less_than"]
+            row_great_than = params["row_great_than"]
+            col_less_than = params["col_less_than"]
+            col_great_than = params["col_great_than"]
+        else:
+            raise ValueError("Invalid flower mode")
+        params = {
+            "flower_mode": flower_mode,
+            "random_dropout": random_dropout,
+            "row_less_than": row_less_than,
+            "row_great_than": row_great_than,
+            "col_less_than": col_less_than,
+            "col_great_than": col_great_than,
+            "bezier_middle_x": params["bezier_middle_x"],
+            "bezier_middle_y": params["bezier_middle_y"],
+            "bezier_end_x": params["bezier_end_x"],
+            "bezier_end_y": params["bezier_end_y"],
+            "flower_num_core_segments": int(params["flower_num_core_segments"]),
+            "flower_num_core_rings": int(params["flower_num_core_rings"]),
+            "flower_radius": uniform(0.02, 0.05),
+            "flower_core_shape_x": uniform(0.8, 1.2),
+            "flower_core_shape_y": uniform(0.8, 1.2),
+            "flower_core_shape_z": uniform(0.5, 0.8),
+            "pedal_stem_mid_point_x": normal(0.0, 0.05),
+            "pedal_stem_mid_point_y": normal(0.0, 0.05),
+            "pedal_stem_radius": uniform(0.02, 0.045),
+            "pedal_stem_top_radius": uniform(0.005, 0.008),
+            "transform_scale": params["transform_scale"],
+            "stem_map_range": params["stem_map_range"],
+            "stem_radius": params["stem_radius"],
+        }
+        self.params.update(params)
+        
+
+    def create_asset(self, **params):
+        bpy.ops.mesh.primitive_plane_add(
+            size=1,
+            enter_editmode=False,
+            align="WORLD",
+            location=(0, 0, 0),
+            scale=(1, 1, 1),
+        )
+        obj = bpy.context.active_object
+
+        surface.add_geomod(
+            obj,
+            geometry_dandelion_nodes,
+            apply=True,
+            attributes=[],
+            input_kwargs=self.params,
+        )
+        tag_object(obj, "dandelion")
+        return obj
+
+
+class DandelionSeedFactory(AssetFactory):
+    def __init__(self, factory_seed, coarse=False):
+        super(DandelionSeedFactory, self).__init__(factory_seed, coarse=coarse)
+
+    def create_asset(self, **params):
+        bpy.ops.mesh.primitive_plane_add(
+            size=1,
+            enter_editmode=False,
+            align="WORLD",
+            location=(0, 0, 0),
+            scale=(1, 1, 1),
+        )
+        obj = bpy.context.active_object
+
+        surface.add_geomod(
+            obj,
+            geometry_dandelion_seed_nodes,
+            apply=True,
+            attributes=[],
+            input_kwargs=params,
+        )
+        tag_object(obj, "seed")
+        return obj
+
+
+if __name__ == "__main__":
+    f = DandelionSeedFactory(0)
+    obj = f.create_asset()

core/assets/flower.py

@@ -0,0 +1,1002 @@
+# Copyright (C) 2023, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Alexander Raistrick, Alejandro Newell
+
+
+# Code generated using version v2.0.1 of the node_transpiler
+import bpy
+import numpy as np
+from numpy.random import normal, uniform
+
+import infinigen
+from infinigen.core import surface
+from infinigen.core.nodes import node_utils
+from infinigen.core.nodes.node_wrangler import Nodes
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.tagging import tag_nodegroup, tag_object
+from infinigen.core.util import blender as butil
+from infinigen.core.util import color
+from infinigen.core.util.math import FixedSeed, dict_lerp
+
+
+@node_utils.to_nodegroup("nodegroup_polar_to_cart_old", singleton=True)
+def nodegroup_polar_to_cart_old(nw):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketVector", "Addend", (0.0, 0.0, 0.0)),
+            ("NodeSocketFloat", "Value", 0.5),
+            ("NodeSocketVector", "Vector", (0.0, 0.0, 0.0)),
+        ],
+    )
+
+    cosine = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Value"]},
+        attrs={"operation": "COSINE"},
+    )
+
+    sine = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Value"]},
+        attrs={"operation": "SINE"},
+    )
+
+    combine_xyz_4 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Y": cosine, "Z": sine})
+
+    multiply_add = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={
+            0: group_input.outputs["Vector"],
+            1: combine_xyz_4,
+            2: group_input.outputs["Addend"],
+        },
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Vector": multiply_add.outputs["Vector"]}
+    )
+
+
+@node_utils.to_nodegroup("nodegroup_follow_curve", singleton=True)
+def nodegroup_follow_curve(nw):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Geometry", None),
+            ("NodeSocketGeometry", "Curve", None),
+            ("NodeSocketFloat", "Curve Min", 0.5),
+            ("NodeSocketFloat", "Curve Max", 1.0),
+        ],
+    )
+
+    position = nw.new_node(Nodes.InputPosition)
+
+    capture_attribute = nw.new_node(
+        Nodes.CaptureAttribute,
+        input_kwargs={"Geometry": group_input.outputs["Geometry"], 1: position},
+        attrs={"data_type": "FLOAT_VECTOR"},
+    )
+
+    separate_xyz = nw.new_node(
+        Nodes.SeparateXYZ,
+        input_kwargs={"Vector": capture_attribute.outputs["Attribute"]},
+    )
+
+    attribute_statistic = nw.new_node(
+        Nodes.AttributeStatistic,
+        input_kwargs={
+            "Geometry": capture_attribute.outputs["Geometry"],
+            2: separate_xyz.outputs["Z"],
+        },
+    )
+
+    map_range = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={
+            "Value": separate_xyz.outputs["Z"],
+            1: attribute_statistic.outputs["Min"],
+            2: attribute_statistic.outputs["Max"],
+            3: group_input.outputs["Curve Min"],
+            4: group_input.outputs["Curve Max"],
+        },
+    )
+
+    curve_length = nw.new_node(
+        Nodes.CurveLength, input_kwargs={"Curve": group_input.outputs["Curve"]}
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: map_range.outputs["Result"], 1: curve_length},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    sample_curve = nw.new_node(
+        Nodes.SampleCurve,
+        input_kwargs={"Curves": group_input.outputs["Curve"], "Length": multiply},
+        attrs={"mode": "LENGTH"},
+    )
+
+    cross_product = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={
+            0: sample_curve.outputs["Tangent"],
+            1: sample_curve.outputs["Normal"],
+        },
+        attrs={"operation": "CROSS_PRODUCT"},
+    )
+
+    scale = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={
+            0: cross_product.outputs["Vector"],
+            "Scale": separate_xyz.outputs["X"],
+        },
+        attrs={"operation": "SCALE"},
+    )
+
+    scale_1 = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={
+            0: sample_curve.outputs["Normal"],
+            "Scale": separate_xyz.outputs["Y"],
+        },
+        attrs={"operation": "SCALE"},
+    )
+
+    add = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={0: scale.outputs["Vector"], 1: scale_1.outputs["Vector"]},
+    )
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={
+            "Geometry": capture_attribute.outputs["Geometry"],
+            "Position": sample_curve.outputs["Position"],
+            "Offset": add.outputs["Vector"],
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": set_position}
+    )
+
+
+@node_utils.to_nodegroup("nodegroup_norm_index", singleton=True)
+def nodegroup_norm_index(nw):
+    index = nw.new_node(Nodes.Index)
+
+    group_input = nw.new_node(
+        Nodes.GroupInput, expose_input=[("NodeSocketInt", "Count", 0)]
+    )
+
+    divide = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: index, 1: group_input.outputs["Count"]},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    group_output = nw.new_node(Nodes.GroupOutput, input_kwargs={"T": divide})
+
+
+@node_utils.to_nodegroup("nodegroup_flower_petal", singleton=True)
+def nodegroup_flower_petal(nw):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Geometry", None),
+            ("NodeSocketFloat", "Length", 0.2),
+            ("NodeSocketFloat", "Point", 1.0),
+            ("NodeSocketFloat", "Point height", 0.5),
+            ("NodeSocketFloat", "Bevel", 6.8),
+            ("NodeSocketFloat", "Base width", 0.2),
+            ("NodeSocketFloat", "Upper width", 0.3),
+            ("NodeSocketInt", "Resolution H", 8),
+            ("NodeSocketInt", "Resolution V", 4),
+            ("NodeSocketFloat", "Wrinkle", 0.1),
+            ("NodeSocketFloat", "Curl", 0.0),
+        ],
+    )
+
+    multiply_add = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Resolution H"], 1: 2.0, 2: 1.0},
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    grid = nw.new_node(
+        Nodes.MeshGrid,
+        input_kwargs={
+            "Vertices X": group_input.outputs["Resolution V"],
+            "Vertices Y": multiply_add,
+        },
+    )
+
+    position = nw.new_node(Nodes.InputPosition)
+
+    capture_attribute = nw.new_node(
+        Nodes.CaptureAttribute,
+        input_kwargs={"Geometry": grid, 1: position},
+        attrs={"data_type": "FLOAT_VECTOR"},
+    )
+
+    separate_xyz = nw.new_node(
+        Nodes.SeparateXYZ,
+        input_kwargs={"Vector": capture_attribute.outputs["Attribute"]},
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: separate_xyz.outputs["X"], 1: 0.05},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": multiply, "Y": separate_xyz.outputs["Y"]}
+    )
+
+    noise_texture = nw.new_node(
+        Nodes.NoiseTexture,
+        input_kwargs={
+            "Vector": combine_xyz,
+            "Scale": 7.9,
+            "Detail": 0.0,
+            "Distortion": 0.2,
+        },
+        attrs={"noise_dimensions": "2D"},
+    )
+
+    add = nw.new_node(
+        Nodes.Math, input_kwargs={0: noise_texture.outputs["Fac"], 1: -0.5}
+    )
+
+    multiply_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: add, 1: group_input.outputs["Wrinkle"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    separate_xyz_1 = nw.new_node(
+        Nodes.SeparateXYZ,
+        input_kwargs={"Vector": capture_attribute.outputs["Attribute"]},
+    )
+
+    add_1 = nw.new_node(Nodes.Math, input_kwargs={0: separate_xyz_1.outputs["X"]})
+
+    absolute = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: separate_xyz_1.outputs["Y"]},
+        attrs={"operation": "ABSOLUTE"},
+    )
+
+    multiply_2 = nw.new_node(
+        Nodes.Math, input_kwargs={0: absolute, 1: 2.0}, attrs={"operation": "MULTIPLY"}
+    )
+
+    power = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: multiply_2, 1: group_input.outputs["Bevel"]},
+        attrs={"operation": "POWER"},
+    )
+
+    multiply_add_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: power, 1: -1.0, 2: 1.0},
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    multiply_3 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: add_1, 1: multiply_add_1},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    multiply_add_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: multiply_3,
+            1: group_input.outputs["Upper width"],
+            2: group_input.outputs["Base width"],
+        },
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    multiply_4 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: separate_xyz_1.outputs["Y"], 1: multiply_add_2},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    power_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: absolute, 1: group_input.outputs["Point"]},
+        attrs={"operation": "POWER"},
+    )
+
+    multiply_add_3 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: power_1, 1: -1.0, 2: 1.0},
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    multiply_5 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: multiply_add_3, 1: group_input.outputs["Point height"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    multiply_add_4 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Point height"], 1: -1.0, 2: 1.0},
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    add_2 = nw.new_node(Nodes.Math, input_kwargs={0: multiply_5, 1: multiply_add_4})
+
+    multiply_6 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: add_2, 1: multiply_add_1},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    multiply_7 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: add_1, 1: multiply_6},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={"X": multiply_1, "Y": multiply_4, "Z": multiply_7},
+    )
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={
+            "Geometry": capture_attribute.outputs["Geometry"],
+            "Position": combine_xyz_1,
+        },
+    )
+
+    multiply_8 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Length"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_3 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Y": multiply_8})
+
+    reroute = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Curl"]}
+    )
+
+    group_1 = nw.new_node(
+        nodegroup_polar_to_cart_old().name,
+        input_kwargs={"Addend": combine_xyz_3, "Value": reroute, "Vector": multiply_8},
+    )
+
+    quadratic_bezier = nw.new_node(
+        Nodes.QuadraticBezier,
+        input_kwargs={
+            "Resolution": 8,
+            "Start": (0.0, 0.0, 0.0),
+            "Middle": combine_xyz_3,
+            "End": group_1,
+        },
+    )
+
+    group = nw.new_node(
+        nodegroup_follow_curve().name,
+        input_kwargs={
+            "Geometry": set_position,
+            "Curve": quadratic_bezier,
+            "Curve Min": 0.0,
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": tag_nodegroup(nw, group, "petal")}
+    )
+
+
+@node_utils.to_nodegroup("nodegroup_phyllo_points", singleton=True)
+def nodegroup_phyllo_points(nw):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketInt", "Count", 50),
+            ("NodeSocketFloat", "Min Radius", 0.0),
+            ("NodeSocketFloat", "Max Radius", 2.0),
+            ("NodeSocketFloat", "Radius exp", 0.5),
+            ("NodeSocketFloat", "Min angle", -0.5236),
+            ("NodeSocketFloat", "Max angle", 0.7854),
+            ("NodeSocketFloat", "Min z", 0.0),
+            ("NodeSocketFloat", "Max z", 1.0),
+            ("NodeSocketFloat", "Clamp z", 1.0),
+            ("NodeSocketFloat", "Yaw offset", -1.5708),
+        ],
+    )
+
+    mesh_line = nw.new_node(
+        Nodes.MeshLine, input_kwargs={"Count": group_input.outputs["Count"]}
+    )
+
+    mesh_to_points = nw.new_node(Nodes.MeshToPoints, input_kwargs={"Mesh": mesh_line})
+
+    position = nw.new_node(Nodes.InputPosition)
+
+    capture_attribute = nw.new_node(
+        Nodes.CaptureAttribute,
+        input_kwargs={"Geometry": mesh_to_points, 1: position},
+        attrs={"data_type": "FLOAT_VECTOR"},
+    )
+
+    index = nw.new_node(Nodes.Index)
+
+    cosine = nw.new_node(
+        Nodes.Math, input_kwargs={0: index}, attrs={"operation": "COSINE"}
+    )
+
+    sine = nw.new_node(Nodes.Math, input_kwargs={0: index}, attrs={"operation": "SINE"})
+
+    combine_xyz = nw.new_node(Nodes.CombineXYZ, input_kwargs={"X": cosine, "Y": sine})
+
+    divide = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: index, 1: group_input.outputs["Count"]},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    power = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: divide, 1: group_input.outputs["Radius exp"]},
+        attrs={"operation": "POWER"},
+    )
+
+    map_range = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={
+            "Value": power,
+            3: group_input.outputs["Min Radius"],
+            4: group_input.outputs["Max Radius"],
+        },
+    )
+
+    multiply = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={0: combine_xyz, 1: map_range.outputs["Result"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    separate_xyz = nw.new_node(
+        Nodes.SeparateXYZ, input_kwargs={"Vector": multiply.outputs["Vector"]}
+    )
+
+    map_range_2 = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={
+            "Value": divide,
+            2: group_input.outputs["Clamp z"],
+            3: group_input.outputs["Min z"],
+            4: group_input.outputs["Max z"],
+        },
+    )
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={
+            "X": separate_xyz.outputs["X"],
+            "Y": separate_xyz.outputs["Y"],
+            "Z": map_range_2.outputs["Result"],
+        },
+    )
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={
+            "Geometry": capture_attribute.outputs["Geometry"],
+            "Position": combine_xyz_1,
+        },
+    )
+
+    map_range_3 = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={
+            "Value": divide,
+            3: group_input.outputs["Min angle"],
+            4: group_input.outputs["Max angle"],
+        },
+    )
+
+    random_value = nw.new_node(Nodes.RandomValue, input_kwargs={2: -0.1, 3: 0.1})
+
+    add = nw.new_node(
+        Nodes.Math, input_kwargs={0: index, 1: group_input.outputs["Yaw offset"]}
+    )
+
+    combine_xyz_2 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={
+            "X": map_range_3.outputs["Result"],
+            "Y": random_value.outputs[1],
+            "Z": add,
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Points": set_position, "Rotation": combine_xyz_2},
+    )
+
+
+@node_utils.to_nodegroup("nodegroup_plant_seed", singleton=True)
+def nodegroup_plant_seed(nw):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketVector", "Dimensions", (0.0, 0.0, 0.0)),
+            ("NodeSocketIntUnsigned", "U", 4),
+            ("NodeSocketInt", "V", 8),
+        ],
+    )
+
+    separate_xyz = nw.new_node(
+        Nodes.SeparateXYZ, input_kwargs={"Vector": group_input.outputs["Dimensions"]}
+    )
+
+    combine_xyz = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"X": separate_xyz.outputs["X"]}
+    )
+
+    multiply_add = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={0: combine_xyz, 1: (0.5, 0.5, 0.5)},
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    quadratic_bezier_1 = nw.new_node(
+        Nodes.QuadraticBezier,
+        input_kwargs={
+            "Resolution": group_input.outputs["U"],
+            "Start": (0.0, 0.0, 0.0),
+            "Middle": multiply_add.outputs["Vector"],
+            "End": combine_xyz,
+        },
+    )
+
+    group = nw.new_node(
+        nodegroup_norm_index().name, input_kwargs={"Count": group_input.outputs["U"]}
+    )
+
+    float_curve = nw.new_node(Nodes.FloatCurve, input_kwargs={"Value": group})
+    node_utils.assign_curve(
+        float_curve.mapping.curves[0], [(0.0, 0.0), (0.3159, 0.4469), (1.0, 0.0156)]
+    )
+
+    map_range = nw.new_node(Nodes.MapRange, input_kwargs={"Value": float_curve, 4: 3.0})
+
+    set_curve_radius = nw.new_node(
+        Nodes.SetCurveRadius,
+        input_kwargs={
+            "Curve": quadratic_bezier_1,
+            "Radius": map_range.outputs["Result"],
+        },
+    )
+
+    curve_circle = nw.new_node(
+        Nodes.CurveCircle,
+        input_kwargs={
+            "Resolution": group_input.outputs["V"],
+            "Radius": separate_xyz.outputs["Y"],
+        },
+    )
+
+    curve_to_mesh = nw.new_node(
+        Nodes.CurveToMesh,
+        input_kwargs={
+            "Curve": set_curve_radius,
+            "Profile Curve": curve_circle.outputs["Curve"],
+            "Fill Caps": True,
+        },
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Mesh": tag_nodegroup(nw, curve_to_mesh, "seed")},
+    )
+
+
+def shader_flower_center(nw):
+    ambient_occlusion = nw.new_node(Nodes.AmbientOcclusion)
+
+    colorramp = nw.new_node(
+        Nodes.ColorRamp, input_kwargs={"Fac": ambient_occlusion.outputs["Color"]}
+    )
+    colorramp.color_ramp.elements.new(1)
+    colorramp.color_ramp.elements[0].position = 0.4841
+    colorramp.color_ramp.elements[0].color = (0.0127, 0.0075, 0.0026, 1.0)
+    colorramp.color_ramp.elements[1].position = 0.8591
+    colorramp.color_ramp.elements[1].color = (0.0848, 0.0066, 0.0007, 1.0)
+    colorramp.color_ramp.elements[2].position = 1.0
+    colorramp.color_ramp.elements[2].color = (1.0, 0.6228, 0.1069, 1.0)
+
+    principled_bsdf = nw.new_node(
+        Nodes.PrincipledBSDF, input_kwargs={"Base Color": colorramp.outputs["Color"]}
+    )
+
+    material_output = nw.new_node(
+        Nodes.MaterialOutput, input_kwargs={"Surface": principled_bsdf}
+    )
+
+
+def shader_petal(nw):
+    translucent_color_change = uniform(0.1, 0.6)
+    specular = normal(0.6, 0.1)
+    roughness = normal(0.4, 0.05)
+    translucent_amt = normal(0.3, 0.05)
+
+    petal_color = nw.new_node(Nodes.RGB)
+    petal_color.outputs[0].default_value = color.color_category("petal")
+
+    translucent_color = nw.new_node(
+        Nodes.MixRGB,
+        [translucent_color_change, petal_color, color.color_category("petal")],
+    )
+
+    translucent_bsdf = nw.new_node(
+        Nodes.TranslucentBSDF, input_kwargs={"Color": translucent_color}
+    )
+
+    principled_bsdf = nw.new_node(
+        Nodes.PrincipledBSDF,
+        input_kwargs={
+            "Base Color": petal_color,
+            "Specular": specular,
+            "Roughness": roughness,
+        },
+    )
+
+    mix_shader = nw.new_node(
+        Nodes.MixShader,
+        input_kwargs={"Fac": translucent_amt, 1: principled_bsdf, 2: translucent_bsdf},
+    )
+
+    material_output = nw.new_node(
+        Nodes.MaterialOutput, input_kwargs={"Surface": mix_shader}
+    )
+
+
+def geo_flower(nw, petal_material, center_material):
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Geometry", None),
+            ("NodeSocketFloat", "Center Rad", 0.0),
+            ("NodeSocketVector", "Petal Dims", (0.0, 0.0, 0.0)),
+            ("NodeSocketFloat", "Seed Size", 0.0),
+            ("NodeSocketFloat", "Min Petal Angle", 0.1),
+            ("NodeSocketFloat", "Max Petal Angle", 1.36),
+            ("NodeSocketFloat", "Wrinkle", 0.01),
+            ("NodeSocketFloat", "Curl", 13.89),
+        ],
+    )
+
+    uv_sphere = nw.new_node(
+        Nodes.MeshUVSphere,
+        input_kwargs={
+            "Segments": 8,
+            "Rings": 8,
+            "Radius": group_input.outputs["Center Rad"],
+        },
+    )
+
+    transform = nw.new_node(
+        Nodes.Transform, input_kwargs={"Geometry": uv_sphere, "Scale": (1.0, 1.0, 0.05)}
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Seed Size"], 1: 1.5},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    distribute_points_on_faces = nw.new_node(
+        Nodes.DistributePointsOnFaces,
+        input_kwargs={
+            "Mesh": transform,
+            "Distance Min": multiply,
+            "Density Max": 50000.0,
+        },
+        attrs={"distribute_method": "POISSON"},
+    )
+
+    multiply_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Seed Size"], 1: 10.0},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={"X": multiply_1, "Y": group_input.outputs["Seed Size"]},
+    )
+
+    group_3 = nw.new_node(
+        nodegroup_plant_seed().name,
+        input_kwargs={"Dimensions": combine_xyz, "U": 6, "V": 6},
+    )
+
+    musgrave_texture = nw.new_node(
+        Nodes.MusgraveTexture,
+        input_kwargs={"W": 13.8, "Scale": 2.41},
+        attrs={"musgrave_dimensions": "4D"},
+    )
+
+    map_range = nw.new_node(
+        Nodes.MapRange, input_kwargs={"Value": musgrave_texture, 3: 0.34, 4: 1.21}
+    )
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ,
+        input_kwargs={"X": map_range.outputs["Result"], "Y": 1.0, "Z": 1.0},
+    )
+
+    instance_on_points_1 = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": distribute_points_on_faces.outputs["Points"],
+            "Instance": group_3,
+            "Rotation": (0.0, -1.5708, 0.0541),
+            "Scale": combine_xyz_1,
+        },
+    )
+
+    realize_instances = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": instance_on_points_1}
+    )
+
+    join_geometry_1 = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [realize_instances, transform]}
+    )
+
+    set_material_1 = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={"Geometry": join_geometry_1, "Material": center_material},
+    )
+
+    multiply_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Center Rad"], 1: 6.2832},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    separate_xyz = nw.new_node(
+        Nodes.SeparateXYZ, input_kwargs={"Vector": group_input.outputs["Petal Dims"]}
+    )
+
+    divide = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: multiply_2, 1: separate_xyz.outputs["Y"]},
+        attrs={"operation": "DIVIDE"},
+    )
+
+    multiply_3 = nw.new_node(
+        Nodes.Math, input_kwargs={0: divide, 1: 1.2}, attrs={"operation": "MULTIPLY"}
+    )
+
+    reroute_3 = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Center Rad"]}
+    )
+
+    reroute_1 = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Min Petal Angle"]}
+    )
+
+    reroute = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Max Petal Angle"]}
+    )
+
+    group_1 = nw.new_node(
+        nodegroup_phyllo_points().name,
+        input_kwargs={
+            "Count": multiply_3,
+            "Min Radius": reroute_3,
+            "Max Radius": reroute_3,
+            "Radius exp": 0.0,
+            "Min angle": reroute_1,
+            "Max angle": reroute,
+            "Max z": 0.0,
+        },
+    )
+
+    subtract = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: separate_xyz.outputs["Z"], 1: separate_xyz.outputs["Y"]},
+        attrs={"operation": "SUBTRACT", "use_clamp": True},
+    )
+
+    reroute_2 = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Wrinkle"]}
+    )
+
+    reroute_4 = nw.new_node(
+        Nodes.Reroute, input_kwargs={"Input": group_input.outputs["Curl"]}
+    )
+
+    group = nw.new_node(
+        nodegroup_flower_petal().name,
+        input_kwargs={
+            "Length": separate_xyz.outputs["X"],
+            "Point": 0.56,
+            "Point height": -0.1,
+            "Bevel": 1.83,
+            "Base width": separate_xyz.outputs["Y"],
+            "Upper width": subtract,
+            "Resolution H": 8,
+            "Resolution V": 16,
+            "Wrinkle": reroute_2,
+            "Curl": reroute_4,
+        },
+    )
+
+    instance_on_points = nw.new_node(
+        Nodes.InstanceOnPoints,
+        input_kwargs={
+            "Points": group_1.outputs["Points"],
+            "Instance": group,
+            "Rotation": group_1.outputs["Rotation"],
+        },
+    )
+
+    realize_instances_1 = nw.new_node(
+        Nodes.RealizeInstances, input_kwargs={"Geometry": instance_on_points}
+    )
+
+    noise_texture = nw.new_node(
+        Nodes.NoiseTexture,
+        input_kwargs={"Scale": 3.73, "Detail": 5.41, "Distortion": -1.0},
+    )
+
+    subtract_1 = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={0: noise_texture.outputs["Color"], 1: (0.5, 0.5, 0.5)},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    value = nw.new_node(Nodes.Value)
+    value.outputs[0].default_value = 0.025
+
+    multiply_4 = nw.new_node(
+        Nodes.VectorMath,
+        input_kwargs={0: subtract_1.outputs["Vector"], 1: value},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    set_position = nw.new_node(
+        Nodes.SetPosition,
+        input_kwargs={
+            "Geometry": realize_instances_1,
+            "Offset": multiply_4.outputs["Vector"],
+        },
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={"Geometry": set_position, "Material": petal_material},
+    )
+
+    join_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [set_material_1, set_material]}
+    )
+
+    set_shade_smooth = nw.new_node(
+        Nodes.SetShadeSmooth,
+        input_kwargs={"Geometry": join_geometry, "Shade Smooth": False},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput, input_kwargs={"Geometry": set_shade_smooth}
+    )
+
+
+class FlowerFactory(AssetFactory):
+    def __init__(self, factory_seed, rad=0.15, diversity_fac=0.25):
+        super(FlowerFactory, self).__init__(factory_seed=factory_seed)
+
+        self.get_params_dict()
+
+        self.rad = rad
+        self.diversity_fac = diversity_fac
+
+        with FixedSeed(factory_seed):
+            self.petal_material = surface.shaderfunc_to_material(shader_petal)
+            self.center_material = surface.shaderfunc_to_material(shader_flower_center)
+            #self.species_params = self.get_flower_params(self.rad)
+            self.params = self.get_flower_params(self.rad * normal(1.0, 0.05))
+
+    def get_params_dict(self):
+        self.params_dict = {
+            "overall_rad": ['continuous', (0.7, 1.3)],
+            "pct_inner": ['continuous', (0.05, 0.5)],
+            "base_width": ['continuous', (4, 16)],
+            "top_width": ['continuous', (0.0, 1.6)],
+            "min_angle": ['continuous', (-20, 100)],
+            "max_angle": ['continuous', (-20, 100)],
+            "seed_size": ['continuous', (0.005, 0.03)],
+            "wrinkle": ['continuous', (0.003, 0.02)],
+            "curl": ['continuous', (-120, 120)],
+        }
+
+    @staticmethod
+    def get_flower_params(overall_rad=0.05):
+        pct_inner = uniform(0.05, 0.4)
+        base_width = 2 * np.pi * overall_rad * pct_inner / normal(20, 5)
+        top_width = overall_rad * np.clip(normal(0.7, 0.3), base_width * 1.2, 100)
+
+        min_angle, max_angle = np.deg2rad(np.sort(uniform(-20, 100, 2)))
+
+        return {
+            "Center Rad": overall_rad * pct_inner,
+            "Petal Dims": np.array(
+                [overall_rad * (1 - pct_inner), base_width, top_width], dtype=np.float32
+            ),
+            "Seed Size": uniform(0.005, 0.01),
+            "Min Petal Angle": min_angle,
+            "Max Petal Angle": max_angle,
+            "Wrinkle": uniform(0.003, 0.02),
+            "Curl": np.deg2rad(normal(30, 50)),
+        }
+
+    def update_params(self, params):
+        overall_rad = params['overall_rad']
+        pct_inner = params['pct_inner']
+        base_width = 2 * np.pi * overall_rad * pct_inner / params['base_width']
+        top_width = overall_rad * np.clip(params['top_width'], base_width * 1.2, 100)
+
+        min_angle = np.deg2rad(params['min_angle'])
+        max_angle = np.deg2rad(params['max_angle'])
+        if min_angle > max_angle:
+            min_angle, max_angle = max_angle, min_angle
+
+        parameters = {
+            "Center Rad": overall_rad * pct_inner,
+            "Petal Dims": np.array(
+                [overall_rad * (1 - pct_inner), base_width, top_width], dtype=np.float32
+            ),
+            "Seed Size": params['seed_size'],
+            "Min Petal Angle": min_angle,
+            "Max Petal Angle": max_angle,
+            "Wrinkle": params['wrinkle'],
+            "Curl": np.deg2rad(params['curl']),
+        }
+        self.params.update(parameters)
+        self.petal_material = surface.shaderfunc_to_material(shader_petal)
+        self.center_material = surface.shaderfunc_to_material(shader_flower_center)
+    
+    def fix_unused_params(self, params):
+        return params
+
+    def create_asset(self, **kwargs) -> bpy.types.Object:
+        vert = butil.spawn_vert("flower")
+        mod = surface.add_geomod(
+            vert,
+            geo_flower,
+            input_kwargs={
+                "petal_material": self.petal_material,
+                "center_material": self.center_material,
+            },
+        )
+
+        #inst_params = self.get_flower_params(self.rad * normal(1, 0.05))
+        #params = dict_lerp(self.species_params, inst_params, 0.25)
+        butil.set_geomod_inputs(mod, self.params)
+
+        butil.apply_modifiers(vert, mod=mod)
+
+        vert.rotation_euler.z = uniform(0, 360)
+        tag_object(vert, "flower")
+        return vert

core/assets/table.py

@@ -0,0 +1,493 @@
+# Copyright (C) 2023, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Yiming Zuo
+
+
+import bpy
+from numpy.random import choice, normal, uniform
+
+from infinigen.assets.material_assignments import AssetList
+from infinigen.assets.objects.tables.legs.single_stand import (
+    nodegroup_generate_single_stand,
+)
+from infinigen.assets.objects.tables.legs.square import nodegroup_generate_leg_square
+from infinigen.assets.objects.tables.legs.straight import (
+    nodegroup_generate_leg_straight,
+)
+from infinigen.assets.objects.tables.strechers import nodegroup_strecher
+from infinigen.assets.objects.tables.table_top import nodegroup_generate_table_top
+from infinigen.assets.objects.tables.table_utils import (
+    nodegroup_create_anchors,
+    nodegroup_create_legs_and_strechers,
+)
+from infinigen.core import surface, tagging
+from infinigen.core import tags as t
+from infinigen.core.nodes import node_utils
+
+# from infinigen.assets.materials import metal, metal_shader_list
+# from infinigen.assets.materials.fabrics import fabric
+from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.surface import NoApply
+from infinigen.core.util.math import FixedSeed
+
+
+@node_utils.to_nodegroup(
+    "geometry_create_legs", singleton=False, type="GeometryNodeTree"
+)
+def geometry_create_legs(nw: NodeWrangler, **kwargs):
+    createanchors = nw.new_node(
+        nodegroup_create_anchors().name,
+        input_kwargs={
+            "Profile N-gon": kwargs["Leg Number"],
+            "Profile Width": kwargs["Leg Placement Top Relative Scale"]
+            * kwargs["Top Profile Width"],
+            "Profile Aspect Ratio": kwargs["Top Profile Aspect Ratio"],
+        },
+    )
+
+    if kwargs["Leg Style"] == "single_stand":
+        leg = nw.new_node(
+            nodegroup_generate_single_stand(**kwargs).name,
+            input_kwargs={
+                "Leg Height": kwargs["Leg Height"],
+                "Leg Diameter": kwargs["Leg Diameter"],
+                "Resolution": 64,
+            },
+        )
+
+        leg = nw.new_node(
+            nodegroup_create_legs_and_strechers().name,
+            input_kwargs={
+                "Anchors": createanchors,
+                "Keep Legs": True,
+                "Leg Instance": leg,
+                "Table Height": kwargs["Top Height"],
+                "Leg Bottom Relative Scale": kwargs[
+                    "Leg Placement Bottom Relative Scale"
+                ],
+                "Align Leg X rot": True,
+            },
+        )
+
+    elif kwargs["Leg Style"] == "straight":
+        leg = nw.new_node(
+            nodegroup_generate_leg_straight(**kwargs).name,
+            input_kwargs={
+                "Leg Height": kwargs["Leg Height"],
+                "Leg Diameter": kwargs["Leg Diameter"],
+                "Resolution": 32,
+                "N-gon": kwargs["Leg NGon"],
+                "Fillet Ratio": 0.1,
+            },
+        )
+
+        strecher = nw.new_node(
+            nodegroup_strecher().name,
+            input_kwargs={"Profile Width": kwargs["Leg Diameter"] * 0.5},
+        )
+
+        leg = nw.new_node(
+            nodegroup_create_legs_and_strechers().name,
+            input_kwargs={
+                "Anchors": createanchors,
+                "Keep Legs": True,
+                "Leg Instance": leg,
+                "Table Height": kwargs["Top Height"],
+                "Strecher Instance": strecher,
+                "Strecher Index Increment": kwargs["Strecher Increament"],
+                "Strecher Relative Position": kwargs["Strecher Relative Pos"],
+                "Leg Bottom Relative Scale": kwargs[
+                    "Leg Placement Bottom Relative Scale"
+                ],
+                "Align Leg X rot": True,
+            },
+        )
+
+    elif kwargs["Leg Style"] == "square":
+        leg = nw.new_node(
+            nodegroup_generate_leg_square(**kwargs).name,
+            input_kwargs={
+                "Height": kwargs["Leg Height"],
+                "Width": 0.707
+                * kwargs["Leg Placement Top Relative Scale"]
+                * kwargs["Top Profile Width"]
+                * kwargs["Top Profile Aspect Ratio"],
+                "Has Bottom Connector": (kwargs["Strecher Increament"] > 0),
+                "Profile Width": kwargs["Leg Diameter"],
+            },
+        )
+
+        leg = nw.new_node(
+            nodegroup_create_legs_and_strechers().name,
+            input_kwargs={
+                "Anchors": createanchors,
+                "Keep Legs": True,
+                "Leg Instance": leg,
+                "Table Height": kwargs["Top Height"],
+                "Leg Bottom Relative Scale": kwargs[
+                    "Leg Placement Bottom Relative Scale"
+                ],
+                "Align Leg X rot": True,
+            },
+        )
+
+    else:
+        raise NotImplementedError
+
+    leg = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={"Geometry": leg, "Material": kwargs["LegMaterial"]},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": leg},
+        attrs={"is_active_output": True},
+    )
+
+
+def geometry_assemble_table(nw: NodeWrangler, **kwargs):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    generatetabletop = nw.new_node(
+        nodegroup_generate_table_top().name,
+        input_kwargs={
+            "Thickness": kwargs["Top Thickness"],
+            "N-gon": kwargs["Top Profile N-gon"],
+            "Profile Width": kwargs["Top Profile Width"],
+            "Aspect Ratio": kwargs["Top Profile Aspect Ratio"],
+            "Fillet Ratio": kwargs["Top Profile Fillet Ratio"],
+            "Fillet Radius Vertical": kwargs["Top Vertical Fillet Ratio"],
+        },
+    )
+
+    tabletop_instance = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": generatetabletop,
+            "Translation": (0.0000, 0.0000, kwargs["Top Height"]),
+        },
+    )
+
+    tabletop_instance = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={"Geometry": tabletop_instance, "Material": kwargs["TopMaterial"]},
+    )
+
+    legs = nw.new_node(geometry_create_legs(**kwargs).name)
+
+    join_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [tabletop_instance, legs]}
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": join_geometry},
+        attrs={"is_active_output": True},
+    )
+
+
+class TableDiningFactory(AssetFactory):
+    def __init__(self, factory_seed, coarse=False, dimensions=None):
+        super(TableDiningFactory, self).__init__(factory_seed, coarse=coarse)
+
+        self.dimensions = dimensions
+        self.get_params_dict()
+        self.leg_styles = ["single_stand", "square", "straight"]
+
+        with FixedSeed(factory_seed):
+            self.params = self.sample_parameters(dimensions)
+
+            # self.clothes_scatter = ClothesCover(factory_fn=blanket.BlanketFactory, width=log_uniform(.8, 1.2),
+            #                                     size=uniform(.8, 1.2)) if uniform() < .3 else NoApply()
+            self.clothes_scatter = NoApply()
+            self.material_params, self.scratch, self.edge_wear = (
+                self.get_material_params()
+            )
+        self.params.update(self.material_params)
+    
+    def get_params_dict(self):
+        # list all the parameters (key:name, value: [type, range]) used in this generator
+        self.params_dict = {
+            "ngon": ["discrete", (4, 36)],
+            "dimension_x": ["continuous", (0.9, 2.2)],
+            "dimension_y": ["continuous", (0.9, 2.2)],
+            "dimension_z": ["continuous", (0.5, 0.9)],
+            "leg_style": ["discrete", (0, 1, 2)],
+            "leg_number": ["discrete", (1, 2, 4)],
+            "leg_ngon": ["discrete", (4, 12)],
+            "leg_diameter": ["continuous", (0, 1)],
+            "leg_height": ["continuous", (0.6, 2.0)],
+            "leg_curve_ctrl_pts0": ["continuous", (0, 1)],
+            "leg_curve_ctrl_pts1": ["continuous", (0, 1)],
+            "leg_curve_ctrl_pts2": ["continuous", (0, 1)],
+            "top_scale": ["continuous", (0.6, 0.8)], # leg start point relative position
+            "bottom_scale": ["continuous", (0.9, 1.3)], # leg end point relative position
+            "top_thickness": ["continuous", (0.02, 0.1)],
+            "top_profile_fillet_ratio": ["continuous", (-0.6, 0.6)], # table corner round / square
+            "top_vertical_fillet_ratio": ["continuous", (0.0, 0.2)], # table corner round / square
+            "strecher_relative_pos": ["continuous", (0.15, 0.8)], 
+            "strecher_increament": ["discrete", (0, 1, 2)],
+        }
+
+
+    def get_material_params(self):
+        material_assignments = AssetList["TableDiningFactory"]()
+        params = {
+            "TopMaterial": material_assignments["top"].assign_material(),
+            "LegMaterial": material_assignments["leg"].assign_material(),
+        }
+        wrapped_params = {
+            k: surface.shaderfunc_to_material(v) for k, v in params.items()
+        }
+
+        scratch_prob, edge_wear_prob = material_assignments["wear_tear_prob"]
+        scratch, edge_wear = material_assignments["wear_tear"]
+
+        is_scratch = uniform() < scratch_prob
+        is_edge_wear = uniform() < edge_wear_prob
+        if not is_scratch:
+            scratch = None
+
+        if not is_edge_wear:
+            edge_wear = None
+
+        return wrapped_params, scratch, edge_wear
+
+    @staticmethod
+    def sample_parameters(dimensions):
+        # not used in DI-PCG
+        if dimensions is None:
+            width = uniform(0.91, 1.16)
+
+            if uniform() < 0.7:
+                # oblong
+                length = uniform(1.4, 2.8)
+            else:
+                # approx square
+                length = width * normal(1, 0.1)
+
+            dimensions = (length, width, uniform(0.65, 0.85))
+
+        # all in meters
+        x, y, z = dimensions
+
+        NGon = 4
+
+        leg_style = choice(["straight", "single_stand", "square"], p=[0.5, 0.1, 0.4])
+        # leg_style = choice(['straight'])
+
+        if leg_style == "single_stand":
+            leg_number = 2
+            leg_diameter = uniform(0.22 * x, 0.28 * x)
+
+            leg_curve_ctrl_pts = [
+                (0.0, uniform(0.1, 0.2)),
+                (0.5, uniform(0.1, 0.2)),
+                (0.9, uniform(0.2, 0.3)),
+                (1.0, 1.0),
+            ]
+
+            top_scale = uniform(0.6, 0.7)
+            bottom_scale = 1.0
+
+        elif leg_style == "square":
+            leg_number = 2
+            leg_diameter = uniform(0.07, 0.10)
+
+            leg_curve_ctrl_pts = None
+
+            top_scale = 0.8
+            bottom_scale = 1.0
+
+        elif leg_style == "straight":
+            leg_diameter = uniform(0.05, 0.07)
+
+            leg_number = 4
+
+            leg_curve_ctrl_pts = [
+                (0.0, 1.0),
+                (0.4, uniform(0.85, 0.95)),
+                (1.0, uniform(0.4, 0.6)),
+            ]
+
+            top_scale = 0.8
+            bottom_scale = uniform(1.0, 1.2)
+
+        else:
+            raise NotImplementedError
+
+        top_thickness = uniform(0.03, 0.06)
+
+        parameters = {
+            "Top Profile N-gon": NGon,
+            "Top Profile Width": 1.414 * x,
+            "Top Profile Aspect Ratio": y / x,
+            "Top Profile Fillet Ratio": uniform(0.0, 0.02),
+            "Top Thickness": top_thickness,
+            "Top Vertical Fillet Ratio": uniform(0.1, 0.3),
+            # 'Top Material': choice(['marble', 'tiled_wood', 'metal', 'fabric'], p=[.3, .3, .2, .2]),
+            "Height": z,
+            "Top Height": z - top_thickness,
+            "Leg Number": leg_number,
+            "Leg Style": leg_style,
+            "Leg NGon": 4,
+            "Leg Placement Top Relative Scale": top_scale,
+            "Leg Placement Bottom Relative Scale": bottom_scale,
+            "Leg Height": 1.0,
+            "Leg Diameter": leg_diameter,
+            "Leg Curve Control Points": leg_curve_ctrl_pts,
+            # 'Leg Material': choice(['metal', 'wood', 'glass', 'plastic']),
+            "Strecher Relative Pos": uniform(0.2, 0.6),
+            "Strecher Increament": choice([0, 1, 2]),
+        }
+
+        return parameters
+
+    def fix_unused_params(self, params):
+        if params['leg_style'] == 0:
+            # single stand only allow 1 or 2 legs
+            if params['leg_number'] == 4:
+                params['leg_number'] = 2
+            params['bottom_scale'] = 1.1
+            params['strecher_increament'] = 1
+        elif params['leg_style'] == 1:
+            params['leg_number'] = 2
+            params['leg_curve_ctrl_pts0'] = 0.5
+            params['leg_curve_ctrl_pts1'] = 0.5
+            params['leg_curve_ctrl_pts2'] = 0.5
+            params['bottom_scale'] = 1.1
+            params['top_scale'] = 0.8
+            params['strecher_increament'] = 1
+        elif params['leg_style'] == 2:
+            params['leg_number'] = 4
+            params['leg_curve_ctrl_pts0'] = 0.5
+            params['top_scale'] = 0.8
+        if params['ngon'] == 36:
+            params['top_profile_fillet_ratio'] = 0.0
+            params['top_vertical_fillet_ratio'] = 0.0
+        return params
+
+    def update_params(self, params):
+        x, y, z = params["dimension_x"], params["dimension_y"], params["dimension_z"]
+        NGon = params['ngon']
+
+        leg_style = self.leg_styles[int(params['leg_style'])]
+        
+        if leg_style == "single_stand":
+            leg_number = params['leg_number']
+            if leg_number == 4:
+                leg_number = 2
+            leg_diameter = (0.2 + 0.2 * params['leg_diameter']) * x
+            leg_curve_ctrl_pts = [
+                (0.0, 0.1 + 0.8 * params['leg_curve_ctrl_pts0']),
+                (0.5, 0.1 + 0.8 * params['leg_curve_ctrl_pts1']),
+                (0.9, 0.2 + 0.8 * params['leg_curve_ctrl_pts2']),
+                (1.0, 1.0),
+            ]
+            top_scale = params['top_scale']
+            bottom_scale = 1.0
+            strecher_increament = 1
+
+        elif leg_style == "square":
+            leg_number = 2
+            leg_diameter = 0.05 + 0.2 * params['leg_diameter']
+            leg_curve_ctrl_pts = None
+            top_scale = 0.8
+            bottom_scale = 1.0
+            strecher_increament = 1
+
+        elif leg_style == "straight":
+            leg_diameter = 0.05 + 0.2 * params['leg_diameter']
+            leg_number = 4
+            leg_curve_ctrl_pts = [
+                (0.0, 1.0),
+                (0.4, 0.5 + 0.5 * params['leg_curve_ctrl_pts1']),
+                (1.0, 0.3 + 0.5 * params['leg_curve_ctrl_pts2'])
+            ]
+            top_scale = 0.8
+            bottom_scale = params['bottom_scale']
+            strecher_increament = params["strecher_increament"]
+        else:
+            raise NotImplementedError
+        
+        if params['ngon'] == 36:
+            top_profile_fillet_ratio = 0.0
+            top_vertical_fillet_ratio = 0.0
+        else:
+            top_profile_fillet_ratio = params['top_profile_fillet_ratio']
+            top_vertical_fillet_ratio = params['top_vertical_fillet_ratio']
+
+        top_thickness = params['top_thickness']
+        parameters = {
+            "Top Profile N-gon": NGon,
+            "Top Profile Width": 1.414 * x,
+            "Top Profile Aspect Ratio": y / x,
+            "Top Profile Fillet Ratio": top_profile_fillet_ratio,
+            "Top Thickness": top_thickness,
+            "Top Vertical Fillet Ratio": top_vertical_fillet_ratio,
+            "Height": z,
+            "Top Height": z - top_thickness,
+            "Leg Number": leg_number,
+            "Leg Style": leg_style,
+            "Leg NGon": params['leg_ngon'],
+            "Leg Placement Top Relative Scale": top_scale,
+            "Leg Placement Bottom Relative Scale": bottom_scale,
+            "Leg Height": params['leg_height'],
+            "Leg Diameter": leg_diameter,
+            "Leg Curve Control Points": leg_curve_ctrl_pts,
+            "Strecher Relative Pos": params["strecher_relative_pos"],
+            "Strecher Increament": strecher_increament,
+        }
+        self.params.update(parameters)
+        self.clothes_scatter = NoApply()
+        self.material_params, self.scratch, self.edge_wear = (
+            self.get_material_params()
+        )
+        self.params.update(self.material_params)
+    
+    def create_asset(self, **params):
+        bpy.ops.mesh.primitive_plane_add(
+            size=2,
+            enter_editmode=False,
+            align="WORLD",
+            location=(0, 0, 0),
+            scale=(1, 1, 1),
+        )
+        obj = bpy.context.active_object
+
+        # surface.add_geomod(obj, geometry_assemble_table, apply=False, input_kwargs=self.params)
+        surface.add_geomod(
+            obj, geometry_assemble_table, apply=True, input_kwargs=self.params
+        )
+        tagging.tag_system.relabel_obj(obj)
+        assert tagging.tagged_face_mask(obj, {t.Subpart.SupportSurface}).sum() != 0
+
+        return obj
+
+    def finalize_assets(self, assets):
+        if self.scratch:
+            self.scratch.apply(assets)
+        if self.edge_wear:
+            self.edge_wear.apply(assets)
+
+    # def finalize_assets(self, assets):
+    #    self.clothes_scatter.apply(assets)
+
+
+class SideTableFactory(TableDiningFactory):
+    def __init__(self, factory_seed, coarse=False, dimensions=None):
+        if dimensions is None:
+            w = 0.55 * normal(1, 0.05)
+            h = 0.95 * w * normal(1, 0.05)
+            dimensions = (w, w, h)
+        super().__init__(factory_seed, coarse=coarse, dimensions=dimensions)
+
+
+class CoffeeTableFactory(TableDiningFactory):
+    def __init__(self, factory_seed, coarse=False, dimensions=None):
+        if dimensions is None:
+            dimensions = (uniform(1, 1.5), uniform(0.6, 0.9), uniform(0.4, 0.5))
+        super().__init__(factory_seed, coarse=coarse, dimensions=dimensions)

core/assets/vase.py

@@ -0,0 +1,486 @@
+# Copyright (C) 2023, Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Yiming Zuo
+
+import bpy
+import numpy as np
+from numpy.random import choice, randint, uniform
+
+import infinigen
+import infinigen.core.util.blender as butil
+from infinigen.assets.material_assignments import AssetList
+from infinigen.assets.objects.table_decorations.utils import (
+    nodegroup_lofting,
+    nodegroup_star_profile,
+)
+from infinigen.core import surface
+from infinigen.core.nodes import node_utils
+from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.util.math import FixedSeed
+
+
+class VaseFactory(AssetFactory):
+    def __init__(self, factory_seed, coarse=False, dimensions=None):
+        super(VaseFactory, self).__init__(factory_seed, coarse=coarse)
+
+        if dimensions is None:
+            z = uniform(0.17, 0.5)
+            x = z * uniform(0.3, 0.6)
+            dimensions = (x, x, z)
+        self.dimensions = dimensions
+        self.get_params_dict()
+
+        with FixedSeed(factory_seed):
+            self.params = self.sample_parameters(dimensions)
+            self.material_params, self.scratch, self.edge_wear = (
+                self.get_material_params()
+            )
+
+        self.params.update(self.material_params)
+    
+    def get_params_dict(self):
+        # list all the parameters (key:name, value: [type, range]) used in this generator
+        self.params_dict = {
+            "dimension_x": ["continuous", (0.05, 0.4)],
+            "dimension_z": ["continuous", (0.2, 0.8)],
+            "neck_scale": ["continuous", (0.15, 0.8)],
+            "profile_inner_radius": ["continuous", (0.8, 1.2)],
+            "profile_star_points": ["discrete", (2,3,4,5,6,7,8,9,10,16,18,20,22,24,26,28,30)],
+            "top_scale": ["continuous", (0.6, 1.4)],
+            "neck_mid_position": ["continuous", (0.5, 1.5)],
+            "neck_position": ["continuous", (-0.2, 0.2)],
+            "shoulder_position": ["continuous", (0.1, 0.8)],
+            "shoulder_thickness": ["continuous", (0.1, 0.3)],
+            "foot_scale": ["continuous", (0.2, 0.8)],
+            "foot_height": ["continuous", (0.01, 0.1)],
+        }
+
+    def get_material_params(self):
+        material_assignments = AssetList["VaseFactory"]()
+        params = {
+            "Material": material_assignments["surface"].assign_material(),
+        }
+        wrapped_params = {
+            k: surface.shaderfunc_to_material(v) for k, v in params.items()
+        }
+
+        scratch_prob, edge_wear_prob = material_assignments["wear_tear_prob"]
+        scratch, edge_wear = material_assignments["wear_tear"]
+
+        is_scratch = uniform() < scratch_prob
+        is_edge_wear = uniform() < edge_wear_prob
+        if not is_scratch:
+            scratch = None
+
+        if not is_edge_wear:
+            edge_wear = None
+
+        return wrapped_params, scratch, edge_wear
+
+    @staticmethod
+    def sample_parameters(dimensions):
+        # all in meters
+        if dimensions is None:
+            z = uniform(0.25, 0.40)
+            x = uniform(0.2, 0.4) * z
+            dimensions = (x, x, z)
+
+        x, y, z = dimensions
+
+        U_resolution = 64
+        V_resolution = 64
+
+        neck_scale = uniform(0.2, 0.8)
+
+        parameters = {
+            "Profile Inner Radius": choice([1.0, uniform(0.8, 1.0)]),
+            "Profile Star Points": randint(16, U_resolution // 2 + 1),
+            "U_resolution": U_resolution,
+            "V_resolution": V_resolution,
+            "Height": z,
+            "Diameter": x,
+            "Top Scale": neck_scale * uniform(0.8, 1.2),
+            "Neck Mid Position": uniform(0.7, 0.95),
+            "Neck Position": 0.5 * neck_scale + 0.5 + uniform(-0.05, 0.05),
+            "Neck Scale": neck_scale,
+            "Shoulder Position": uniform(0.3, 0.7),
+            "Shoulder Thickness": uniform(0.1, 0.25),
+            "Foot Scale": uniform(0.4, 0.6),
+            "Foot Height": uniform(0.01, 0.1),
+        }
+
+        return parameters
+
+    def fix_unused_params(self, params):
+        return params
+    
+    def update_params(self, params):
+        x, y, z = params["dimension_x"], params["dimension_x"], params["dimension_z"]
+        U_resolution = 64
+        V_resolution = 64
+        neck_scale = params["neck_scale"]
+        parameters = {
+            "Profile Inner Radius": np.clip(params["profile_inner_radius"], 0.8, 1.0),
+            "Profile Star Points": params["profile_star_points"],
+            "U_resolution": U_resolution,
+            "V_resolution": V_resolution,
+            "Height": z,
+            "Diameter": x,
+            "Top Scale": neck_scale * params["top_scale"],
+            "Neck Mid Position": params["neck_mid_position"],
+            "Neck Position": 0.5 * neck_scale + 0.5 + params["neck_position"],
+            "Neck Scale": neck_scale,
+            "Shoulder Position": params["shoulder_position"],
+            "Shoulder Thickness": params["shoulder_thickness"],
+            "Foot Scale": params["foot_scale"],
+            "Foot Height": params["foot_height"],
+        }
+        self.params.update(parameters)
+        self.material_params, self.scratch, self.edge_wear = (
+                self.get_material_params()
+        )
+
+        self.params.update(self.material_params)
+
+    def create_asset(self, **params):
+        bpy.ops.mesh.primitive_plane_add(
+            size=2,
+            enter_editmode=False,
+            align="WORLD",
+            location=(0, 0, 0),
+            scale=(1, 1, 1),
+        )
+        obj = bpy.context.active_object
+
+        surface.add_geomod(obj, geometry_vases, apply=True, input_kwargs=self.params)
+        butil.modify_mesh(obj, "SOLIDIFY", apply=True, thickness=0.002)
+        butil.modify_mesh(obj, "SUBSURF", apply=True, levels=2, render_levels=2)
+
+        return obj
+
+    def finalize_assets(self, assets):
+        if self.scratch:
+            self.scratch.apply(assets)
+        if self.edge_wear:
+            self.edge_wear.apply(assets)
+
+
+@node_utils.to_nodegroup(
+    "nodegroup_vase_profile", singleton=False, type="GeometryNodeTree"
+)
+def nodegroup_vase_profile(nw: NodeWrangler):
+    # Code generated using version 2.6.4 of the node_transpiler
+
+    group_input = nw.new_node(
+        Nodes.GroupInput,
+        expose_input=[
+            ("NodeSocketGeometry", "Profile Curve", None),
+            ("NodeSocketFloat", "Height", 0.0000),
+            ("NodeSocketFloat", "Diameter", 0.0000),
+            ("NodeSocketFloat", "Top Scale", 0.0000),
+            ("NodeSocketFloat", "Neck Mid Position", 0.0000),
+            ("NodeSocketFloat", "Neck Position", 0.5000),
+            ("NodeSocketFloat", "Neck Scale", 0.0000),
+            ("NodeSocketFloat", "Shoulder Position", 0.0000),
+            ("NodeSocketFloat", "Shoulder Thickness", 0.0000),
+            ("NodeSocketFloat", "Foot Scale", 0.0000),
+            ("NodeSocketFloat", "Foot Height", 0.0000),
+        ],
+    )
+
+    combine_xyz_1 = nw.new_node(
+        Nodes.CombineXYZ, input_kwargs={"Z": group_input.outputs["Height"]}
+    )
+
+    multiply = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Top Scale"],
+            1: group_input.outputs["Diameter"],
+        },
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    neck_top = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input.outputs["Profile Curve"],
+            "Translation": combine_xyz_1,
+            "Scale": multiply,
+        },
+    )
+
+    multiply_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Height"],
+            1: group_input.outputs["Neck Position"],
+        },
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_1})
+
+    multiply_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Diameter"],
+            1: group_input.outputs["Neck Scale"],
+        },
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    neck = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input.outputs["Profile Curve"],
+            "Translation": combine_xyz,
+            "Scale": multiply_2,
+        },
+    )
+
+    subtract = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: 1.0000, 1: group_input.outputs["Neck Position"]},
+        attrs={"use_clamp": True, "operation": "SUBTRACT"},
+    )
+
+    multiply_add = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: subtract,
+            1: group_input.outputs["Neck Mid Position"],
+            2: group_input.outputs["Neck Position"],
+        },
+        attrs={"operation": "MULTIPLY_ADD"},
+    )
+
+    multiply_3 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: multiply_add, 1: group_input.outputs["Height"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_2 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_3})
+
+    add = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Neck Scale"],
+            1: group_input.outputs["Top Scale"],
+        },
+    )
+
+    divide = nw.new_node(
+        Nodes.Math, input_kwargs={0: add, 1: 2.0000}, attrs={"operation": "DIVIDE"}
+    )
+
+    multiply_4 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: group_input.outputs["Diameter"], 1: divide},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    neck_middle = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input.outputs["Profile Curve"],
+            "Translation": combine_xyz_2,
+            "Scale": multiply_4,
+        },
+    )
+
+    neck_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [neck, neck_middle, neck_top]}
+    )
+
+    map_range = nw.new_node(
+        Nodes.MapRange,
+        input_kwargs={
+            "Value": group_input.outputs["Shoulder Position"],
+            3: group_input.outputs["Foot Height"],
+            4: group_input.outputs["Neck Position"],
+        },
+    )
+
+    subtract_1 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Neck Position"],
+            1: group_input.outputs["Foot Height"],
+        },
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    multiply_5 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: subtract_1, 1: group_input.outputs["Shoulder Thickness"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    add_1 = nw.new_node(
+        Nodes.Math, input_kwargs={0: map_range.outputs["Result"], 1: multiply_5}
+    )
+
+    minimum = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: add_1, 1: group_input.outputs["Neck Position"]},
+        attrs={"operation": "MINIMUM"},
+    )
+
+    multiply_6 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: minimum, 1: group_input.outputs["Height"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_3 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_6})
+
+    body_top = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input.outputs["Profile Curve"],
+            "Translation": combine_xyz_3,
+            "Scale": group_input.outputs["Diameter"],
+        },
+    )
+
+    subtract_2 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: map_range.outputs["Result"], 1: multiply_5},
+        attrs={"operation": "SUBTRACT"},
+    )
+
+    maximum = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: subtract_2, 1: group_input.outputs["Foot Height"]},
+        attrs={"operation": "MAXIMUM"},
+    )
+
+    multiply_7 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={0: maximum, 1: group_input.outputs["Height"]},
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_5 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_7})
+
+    body_bottom = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input.outputs["Profile Curve"],
+            "Translation": combine_xyz_5,
+            "Scale": group_input.outputs["Diameter"],
+        },
+    )
+
+    body_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [body_bottom, body_top]}
+    )
+
+    multiply_8 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Foot Height"],
+            1: group_input.outputs["Height"],
+        },
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    combine_xyz_4 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_8})
+
+    multiply_9 = nw.new_node(
+        Nodes.Math,
+        input_kwargs={
+            0: group_input.outputs["Diameter"],
+            1: group_input.outputs["Foot Scale"],
+        },
+        attrs={"operation": "MULTIPLY"},
+    )
+
+    foot_top = nw.new_node(
+        Nodes.Transform,
+        input_kwargs={
+            "Geometry": group_input,
+            "Translation": combine_xyz_4,
+            "Scale": multiply_9,
+        },
+    )
+
+    foot_bottom = nw.new_node(
+        Nodes.Transform, input_kwargs={"Geometry": group_input, "Scale": multiply_9}
+    )
+
+    foot_geometry = nw.new_node(
+        Nodes.JoinGeometry, input_kwargs={"Geometry": [foot_bottom, foot_top]}
+    )
+
+    join_geometry_2 = nw.new_node(
+        Nodes.JoinGeometry,
+        input_kwargs={"Geometry": [foot_geometry, body_geometry, neck_geometry]},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": join_geometry_2},
+        attrs={"is_active_output": True},
+    )
+
+
+def geometry_vases(nw: NodeWrangler, **kwargs):
+    # Code generated using version 2.6.4 of the node_transpiler
+    starprofile = nw.new_node(
+        nodegroup_star_profile().name,
+        input_kwargs={
+            "Resolution": kwargs["U_resolution"],
+            "Points": kwargs["Profile Star Points"],
+            "Inner Radius": kwargs["Profile Inner Radius"],
+        },
+    )
+
+    vaseprofile = nw.new_node(
+        nodegroup_vase_profile().name,
+        input_kwargs={
+            "Profile Curve": starprofile.outputs["Curve"],
+            "Height": kwargs["Height"],
+            "Diameter": kwargs["Diameter"],
+            "Top Scale": kwargs["Top Scale"],
+            "Neck Mid Position": kwargs["Neck Mid Position"],
+            "Neck Position": kwargs["Neck Position"],
+            "Neck Scale": kwargs["Neck Scale"],
+            "Shoulder Position": kwargs["Shoulder Position"],
+            "Shoulder Thickness": kwargs["Shoulder Thickness"],
+            "Foot Scale": kwargs["Foot Scale"],
+            "Foot Height": kwargs["Foot Height"],
+        },
+    )
+
+    lofting = nw.new_node(
+        nodegroup_lofting().name,
+        input_kwargs={
+            "Profile Curves": vaseprofile,
+            "U Resolution": 64,
+            "V Resolution": 64,
+        },
+    )
+
+    delete_geometry = nw.new_node(
+        Nodes.DeleteGeometry,
+        input_kwargs={
+            "Geometry": lofting.outputs["Geometry"],
+            "Selection": lofting.outputs["Top"],
+        },
+    )
+
+    set_material = nw.new_node(
+        Nodes.SetMaterial,
+        input_kwargs={"Geometry": delete_geometry, "Material": kwargs["Material"]},
+    )
+
+    group_output = nw.new_node(
+        Nodes.GroupOutput,
+        input_kwargs={"Geometry": set_material},
+        attrs={"is_active_output": True},
+    )

core/dataset.py

@@ -0,0 +1,40 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+import cv2
+import json
+from core.utils.io import read_list_from_txt
+from core.utils.math_utils import normalize_params
+
+class ImageParamsDataset(Dataset):
+    def __init__(self, data_root, list_file, params_dict_file):
+        self.data_root = data_root
+        self.data_lists = read_list_from_txt(os.path.join(data_root, list_file))
+        self.params_dict = json.load(open(os.path.join(data_root, params_dict_file), 'r'))
+
+    def __len__(self):
+        return len(self.data_lists)
+
+    def __getitem__(self, idx):
+        name = self.data_lists[idx]
+        id = name.split("/")[0]
+        params = json.load(open(os.path.join(self.data_root, id, "params.txt"), 'r'))
+        # normalize the params to [-1, 1] range for training diffusion
+        normalized_params = normalize_params(params, self.params_dict)
+        normalized_params_values = np.array(list(normalized_params.values()))
+        img = cv2.cvtColor(cv2.imread(os.path.join(self.data_root, name)), cv2.COLOR_BGR2RGB)
+        
+        img_feat_name = os.path.join(self.data_root, name.replace(".png", "_dino_token.npy"))
+        if not os.path.exists(img_feat_name):
+            img_feat_file = np.load(os.path.join(self.data_root, name.replace(".png", "_dino_token.npz")))
+            img_feat = img_feat_file['arr_0']
+            img_feat_file.close()
+        else:
+            img_feat = np.load(img_feat_name)
+        img_feat_t = torch.from_numpy(img_feat).float()
+        return torch.from_numpy(normalized_params_values).float(), img_feat_t, img
+
+        

core/diffusion/__init__.py

@@ -0,0 +1,46 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+
+
+def create_diffusion(
+    timestep_respacing,
+    noise_schedule="linear", 
+    use_kl=False,
+    sigma_small=False,
+    predict_xstart=False,
+    learn_sigma=True,
+    rescale_learned_sigmas=False,
+    diffusion_steps=1000
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if timestep_respacing is None or timestep_respacing == "":
+        timestep_respacing = [diffusion_steps]
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=(
+            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
+        ),
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type
+        # rescale_timesteps=rescale_timesteps,
+    )

core/diffusion/diffusion_utils.py

@@ -0,0 +1,88 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+import torch as th
+import numpy as np
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def continuous_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a continuous Gaussian distribution.
+    :param x: the targets
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    normalized_x = centered_x * inv_stdv
+    log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x)
+    return log_probs
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs

core/diffusion/gaussian_diffusion.py

@@ -0,0 +1,873 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+
+import math
+
+import numpy as np
+import torch as th
+import enum
+
+from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
+    betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    warmup_time = int(num_diffusion_timesteps * warmup_frac)
+    betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
+    return betas
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    """
+    This is the deprecated API for creating beta schedules.
+    See get_named_beta_schedule() for the new library of schedules.
+    """
+    if beta_schedule == "quad":
+        betas = (
+            np.linspace(
+                beta_start ** 0.5,
+                beta_end ** 0.5,
+                num_diffusion_timesteps,
+                dtype=np.float64,
+            )
+            ** 2
+        )
+    elif beta_schedule == "linear":
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "warmup10":
+        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1)
+    elif beta_schedule == "warmup50":
+        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5)
+    elif beta_schedule == "const":
+        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
+        betas = 1.0 / np.linspace(
+            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
+        )
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule(
+            "linear",
+            beta_start=scale * 0.0001,
+            beta_end=scale * 0.02,
+            num_diffusion_timesteps=num_diffusion_timesteps,
+        )
+    elif schedule_name == "squaredcos_cap_v2":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+    Original ported from this codebase:
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type
+    ):
+
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        ) if len(self.posterior_variance) > 1 else np.array([])
+
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+        In other words, sample from q(x_t | x_0).
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+            q(x_{t-1} | x_t, x_0)
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, t, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, extra = model_output
+        else:
+            extra = None
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)
+            max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+            # The model_var_values is [-1, 1] for [min_var, max_var].
+            frac = (model_var_values + 1) / 2
+            model_log_variance = frac * max_log + (1 - frac) * min_log
+            model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == ModelMeanType.START_X:
+            pred_xstart = process_xstart(model_output)
+        else:
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+            )
+        model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
+
+        assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+            "extra": extra,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, t, **model_kwargs)
+        new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+        See condition_mean() for details on cond_fn.
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t)
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model.
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+            self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+    ):
+        """
+        Get a term for the variational lower-bound.
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
+        """
+        Compute training losses for a single timestep.
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_output = model(x_t, t, **model_kwargs)
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
+                    x_start=x_start, x_t=x_t, t=t
+                )[0],
+                ModelMeanType.START_X: x_start,
+                ModelMeanType.EPSILON: noise,
+            }[self.model_mean_type]
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+        This term can't be optimized, as it only depends on the encoder.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)

core/diffusion/respace.py

@@ -0,0 +1,129 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(
+            model, self.timestep_map, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        # self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        # if self.rescale_timesteps:
+        #     new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)

core/diffusion/timestep_sampler.py

@@ -0,0 +1,150 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch as th
+import torch.distributed as dist
+
+
+def create_named_schedule_sampler(name, diffusion):
+    """
+    Create a ScheduleSampler from a library of pre-defined samplers.
+    :param name: the name of the sampler.
+    :param diffusion: the diffusion object to sample for.
+    """
+    if name == "uniform":
+        return UniformSampler(diffusion)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(diffusion)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {name}")
+
+
+class ScheduleSampler(ABC):
+    """
+    A distribution over timesteps in the diffusion process, intended to reduce
+    variance of the objective.
+    By default, samplers perform unbiased importance sampling, in which the
+    objective's mean is unchanged.
+    However, subclasses may override sample() to change how the resampled
+    terms are reweighted, allowing for actual changes in the objective.
+    """
+
+    @abstractmethod
+    def weights(self):
+        """
+        Get a numpy array of weights, one per diffusion step.
+        The weights needn't be normalized, but must be positive.
+        """
+
+    def sample(self, batch_size, device):
+        """
+        Importance-sample timesteps for a batch.
+        :param batch_size: the number of timesteps.
+        :param device: the torch device to save to.
+        :return: a tuple (timesteps, weights):
+                 - timesteps: a tensor of timestep indices.
+                 - weights: a tensor of weights to scale the resulting losses.
+        """
+        w = self.weights()
+        p = w / np.sum(w)
+        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
+        indices = th.from_numpy(indices_np).long().to(device)
+        weights_np = 1 / (len(p) * p[indices_np])
+        weights = th.from_numpy(weights_np).float().to(device)
+        return indices, weights
+
+
+class UniformSampler(ScheduleSampler):
+    def __init__(self, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()

core/models.py

@@ -0,0 +1,331 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# GLIDE: https://github.com/openai/glide-text2im
+# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+import xformers.ops
+
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+class MultiHeadCrossAttention(nn.Module):
+    def __init__(self, d_model, num_heads, attn_drop=0., proj_drop=0., **block_kwargs):
+        super(MultiHeadCrossAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+
+        self.q_linear = nn.Linear(d_model, d_model)
+        self.kv_linear = nn.Linear(d_model, d_model*2)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(d_model, d_model)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, cond, mask=None):
+        # query: img tokens; key/value: condition; mask: if padding tokens
+        B, N, C = x.shape
+
+        q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
+        kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
+        k, v = kv.unbind(2)
+        attn_bias = None
+        if mask is not None:
+            attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
+        x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
+        x = x.view(B, -1, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+#################################################################################
+#                                 Core DiT Model                                #
+#################################################################################
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with cross attention for conditioning. Adapted from PixArt implementation.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
+        self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        #self.adaLN_modulation = nn.Sequential(
+        #    nn.SiLU(),
+        #    nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        #)
+        self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5)
+
+    def forward(self, x, y, t, mask=None):
+        B, N, C = x.shape
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1)
+        x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)).reshape(B, N, C)
+        x = x + self.cross_attn(x, y, mask)
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+        self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size ** 0.5)
+        self.out_channels = out_channels
+
+    def forward(self, x, t):
+        shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        input_size=32,
+        in_channels=1,
+        hidden_size=128,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4.0,
+        condition_channels=768,
+        learn_sigma=True,
+    ):
+        super().__init__()
+        self.learn_sigma = learn_sigma
+        self.input_size = input_size
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.num_heads = num_heads
+
+        self.x_embedder = nn.Linear(in_channels, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.t_block = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        self.y_embedder = Mlp(in_features=condition_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=approx_gelu, drop=0)
+        # Will use fixed sin-cos embedding:
+        self.pos_embed = nn.Parameter(torch.zeros(1, input_size, hidden_size), requires_grad=False)
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
+        ])
+        self.final_layer = FinalLayer(hidden_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        grid_1d = np.arange(self.input_size, dtype=np.float32)
+        pos_embed = get_1d_sincos_pos_embed_from_grid(self.pos_embed.shape[-1], grid_1d)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        nn.init.xavier_uniform_(self.x_embedder.weight)
+        nn.init.constant_(self.x_embedder.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.fc1.weight, std=0.02)
+        nn.init.normal_(self.y_embedder.fc2.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.cross_attn.proj.weight, 0)
+            nn.init.constant_(block.cross_attn.proj.bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+    
+    def ckpt_wrapper(self, module):
+        def ckpt_forward(*inputs):
+            outputs = module(*inputs)
+            return outputs
+        return ckpt_forward
+
+    def forward(self, x, t, y):
+        """
+        Forward pass of DiT.
+        x: (N, 1, T) tensor of PCG params
+        t: (N,) tensor of diffusion timesteps
+        y: (N, 1, C) or (N, M, C) tensor of condition image features
+        """
+        x = x.permute(0, 2, 1)
+        x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T is the input token number (params number)
+        t = self.t_embedder(t)                   # (N, D)
+        t0 = self.t_block(t)
+        y = self.y_embedder(y)                   # (N, M, D)
+
+        # mask for batch cross-attention
+        y_lens = [y.shape[1]] * y.shape[0]
+        y = y.view(1, -1, x.shape[-1])
+        for block in self.blocks:
+            x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, y, t0, y_lens)       # (N, T, D)
+        x = self.final_layer(x, t)                # (N, T, out_channels)
+        return x.permute(0, 2, 1)
+
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+#################################################################################
+#                                   DiT Configs                                  #
+#################################################################################
+
+def DiT_S(**kwargs):
+    # 39M
+    return DiT(depth=16, hidden_size=384, num_heads=6, **kwargs)
+
+def DiT_mini(**kwargs):
+    # 7.6M
+    return DiT(depth=12, hidden_size=192, num_heads=6, **kwargs)
+
+def DiT_tiny(**kwargs):
+    # 1.3M
+    return DiT(depth=8, hidden_size=96, num_heads=6, **kwargs)
+
+
+DiT_models = {
+    'DiT_S': DiT_S,
+    'DiT_mini': DiT_mini,
+    'DiT_tiny': DiT_tiny
+}