diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000000000000000000000000000000000000..f6705391aa621ba26a69d8c07229d7eec06799fc --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/infinigen"] + path = third_party/infinigen + url = https://github.com/princeton-vl/infinigen diff --git a/README.md b/README.md index c2e62bae54e55b77bad36cc917e789f5778aa39a..0f8cf1b3f8e91df21f47838b17e3c4589da27b26 100644 --- a/README.md +++ b/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 diff --git a/app.py b/app.py new file mode 100755 index 0000000000000000000000000000000000000000..19b8686a488c861f86e5f730f2bec92ce8c6d437 --- /dev/null +++ b/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_ = ''' +

Official 🤗 Gradio Demo

DI-PCG: Diffusion-based Efficient Inverse Procedural Content Generation for High-quality 3D Asset Creation

+ +**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: GitHub. Techenical report: ArXiv. + +❗️❗️❗️**Important Notes:** +- DI-PCG trains a diffusion model for each procedural generator. Current supported generators are: Chair, Table, Vase, Basket, Flower, Dandelion from Infinigen. +- 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 Github Repo. 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 . +""" +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) \ No newline at end of file diff --git a/configs/demo/basket_demo.yaml b/configs/demo/basket_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..7a212d7966295a0a9c748f2465f313f0ce72ac88 --- /dev/null +++ b/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 diff --git a/configs/demo/chair_demo.yaml b/configs/demo/chair_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..7e02f6d731cc034a6b9f85c85b836a38ddabb02d --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/demo/dandelion_demo.yaml b/configs/demo/dandelion_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..163795170b41aad8efe6983e522b811164576da7 --- /dev/null +++ b/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 diff --git a/configs/demo/flower_demo.yaml b/configs/demo/flower_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..c9651ee8edf7b4f8ba94c7f817a0961f44597736 --- /dev/null +++ b/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 diff --git a/configs/demo/table_demo.yaml b/configs/demo/table_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..902bca857566869622936570b2f1b9ac363f14fb --- /dev/null +++ b/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 diff --git a/configs/demo/vase_demo.yaml b/configs/demo/vase_demo.yaml new file mode 100755 index 0000000000000000000000000000000000000000..7453daeff9313ab32fccdb8e7a18a97ed8ff6f1b --- /dev/null +++ b/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 diff --git a/configs/infinigen/base.gin b/configs/infinigen/base.gin new file mode 100755 index 0000000000000000000000000000000000000000..8b0f01ad2886066fc830ade83bdcc008494072c3 --- /dev/null +++ b/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)} +] diff --git a/configs/test/basket_test.yaml b/configs/test/basket_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..450aa04f338857bb16fa176cb6e0ecf9e1c803f8 --- /dev/null +++ b/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 diff --git a/configs/test/chair_test.yaml b/configs/test/chair_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..043c3647695954c4454e8e9d738ba2f305edbed7 --- /dev/null +++ b/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 diff --git a/configs/test/dandelion_test.yaml b/configs/test/dandelion_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..e200ec822f47a7f4b7dc0876aea435c24ffbd20f --- /dev/null +++ b/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 diff --git a/configs/test/flower_test.yaml b/configs/test/flower_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..9bf648133e3bcd10f94a4297f6479f91cd817b7d --- /dev/null +++ b/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 diff --git a/configs/test/table_test.yaml b/configs/test/table_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..3b79e51957dd825af7ad5d362d4693e405697255 --- /dev/null +++ b/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 diff --git a/configs/test/vase_test.yaml b/configs/test/vase_test.yaml new file mode 100755 index 0000000000000000000000000000000000000000..8233a1cf97e61c38b7d06f99f5b9a4e894eeb648 --- /dev/null +++ b/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 diff --git a/configs/train/basket_train.yaml b/configs/train/basket_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..7a0c3d227691b6ef42be9e14896ca630c040ba97 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train/chair_train.yaml b/configs/train/chair_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..d403578cc2969242b3bc00fde1eeb1f1d8084d27 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train/dandelion_train.yaml b/configs/train/dandelion_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..34457811d0715b8260765d7b5e04ff815d2d3572 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train/flower_train.yaml b/configs/train/flower_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..7d61026d22321b503cad28edc3b920334d5e0b67 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train/table_train.yaml b/configs/train/table_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..2acfaea1feaa9fd05a9da088b62e0f3de79895c5 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/configs/train/vase_train.yaml b/configs/train/vase_train.yaml new file mode 100755 index 0000000000000000000000000000000000000000..f0430c1f07083bccb13667875f2922747f7cefeb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/core/__pycache__/dataset.cpython-310.pyc b/core/__pycache__/dataset.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..d9a571611f24b1584b85bc69cd8dbba31f5e4100 Binary files /dev/null and b/core/__pycache__/dataset.cpython-310.pyc differ diff --git a/core/__pycache__/models.cpython-310.pyc b/core/__pycache__/models.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..d31bf2f6e636d811727db0e14221e0bb4d59d4bb Binary files /dev/null and b/core/__pycache__/models.cpython-310.pyc differ diff --git a/core/assets/__pycache__/basket.cpython-310.pyc b/core/assets/__pycache__/basket.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..f6a16e9c61ac8957d2ac0fba60a3ddfb3c272545 Binary files /dev/null and b/core/assets/__pycache__/basket.cpython-310.pyc differ diff --git a/core/assets/__pycache__/chair.cpython-310.pyc b/core/assets/__pycache__/chair.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..901e043ec4b07e72589079e601d27c7d3c705eb7 Binary files /dev/null and b/core/assets/__pycache__/chair.cpython-310.pyc differ diff --git a/core/assets/__pycache__/dandelion.cpython-310.pyc b/core/assets/__pycache__/dandelion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b793b8c3a120013ad7cde08ee2c98f652bfff85 Binary files /dev/null and b/core/assets/__pycache__/dandelion.cpython-310.pyc differ diff --git a/core/assets/__pycache__/flower.cpython-310.pyc b/core/assets/__pycache__/flower.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..29660555ce82aa37a00c683495295e6482b23ea0 Binary files /dev/null and b/core/assets/__pycache__/flower.cpython-310.pyc differ diff --git a/core/assets/__pycache__/table.cpython-310.pyc b/core/assets/__pycache__/table.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..a5d26ae0ab9caa702bc06af6ac34215038144a77 Binary files /dev/null and b/core/assets/__pycache__/table.cpython-310.pyc differ diff --git a/core/assets/__pycache__/vase.cpython-310.pyc b/core/assets/__pycache__/vase.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..8fd7d788d5f074f8961f4f24874bbffe472db1d8 Binary files /dev/null and b/core/assets/__pycache__/vase.cpython-310.pyc differ diff --git a/core/assets/basket.py b/core/assets/basket.py new file mode 100755 index 0000000000000000000000000000000000000000..30f83ee9192e7567080a415c22aab07aae2cd7e1 --- /dev/null +++ b/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 diff --git a/core/assets/chair.py b/core/assets/chair.py new file mode 100755 index 0000000000000000000000000000000000000000..1e6a051f992e14a4a1eee8993f17aa3f0a65e8bd --- /dev/null +++ b/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 diff --git a/core/assets/dandelion.py b/core/assets/dandelion.py new file mode 100755 index 0000000000000000000000000000000000000000..6151a9a3daca851e509381a2cf82c2e2a68dfc16 --- /dev/null +++ b/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() diff --git a/core/assets/flower.py b/core/assets/flower.py new file mode 100755 index 0000000000000000000000000000000000000000..85ffbe2784316517e8f2940b6970868eea1f4c69 --- /dev/null +++ b/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 diff --git a/core/assets/table.py b/core/assets/table.py new file mode 100755 index 0000000000000000000000000000000000000000..ddbede5d338690e310d389d9ddbf9c292b499b6f --- /dev/null +++ b/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) diff --git a/core/assets/vase.py b/core/assets/vase.py new file mode 100755 index 0000000000000000000000000000000000000000..66972b221ce17dcf4c19818a373cc9986ea4cdd4 --- /dev/null +++ b/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}, + ) diff --git a/core/dataset.py b/core/dataset.py new file mode 100755 index 0000000000000000000000000000000000000000..964a2967c02cae1c473f24231f9d9c383d9f1fd1 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/core/diffusion/__init__.py b/core/diffusion/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..8c536a98da92c4d051458803737661e5ecf974c2 --- /dev/null +++ b/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, + ) diff --git a/core/diffusion/__pycache__/__init__.cpython-310.pyc b/core/diffusion/__pycache__/__init__.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..1785654b2a238072b62a91b0ec01b9e93eb624eb Binary files /dev/null and b/core/diffusion/__pycache__/__init__.cpython-310.pyc differ diff --git a/core/diffusion/__pycache__/diffusion_utils.cpython-310.pyc b/core/diffusion/__pycache__/diffusion_utils.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..6cd4ae363959fa0fecc11889092d62259f12b062 Binary files /dev/null and b/core/diffusion/__pycache__/diffusion_utils.cpython-310.pyc differ diff --git a/core/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc b/core/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..3d363276295cc15b769cfa28e192f0a26fbe9b17 Binary files /dev/null and b/core/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc differ diff --git a/core/diffusion/__pycache__/respace.cpython-310.pyc b/core/diffusion/__pycache__/respace.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..cfb37c68027ce93d0087d61cfb45fded537a6784 Binary files /dev/null and b/core/diffusion/__pycache__/respace.cpython-310.pyc differ diff --git a/core/diffusion/diffusion_utils.py b/core/diffusion/diffusion_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..e493a6a3ecb91e553a53cc7eadee5cc0d1753060 --- /dev/null +++ b/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 diff --git a/core/diffusion/gaussian_diffusion.py b/core/diffusion/gaussian_diffusion.py new file mode 100755 index 0000000000000000000000000000000000000000..ccbcefeca4348e2e627d22723b50492c894e66ba --- /dev/null +++ b/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) diff --git a/core/diffusion/respace.py b/core/diffusion/respace.py new file mode 100755 index 0000000000000000000000000000000000000000..0a2cc0435d1ace54466585db9043b284973d454e --- /dev/null +++ b/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) diff --git a/core/diffusion/timestep_sampler.py b/core/diffusion/timestep_sampler.py new file mode 100755 index 0000000000000000000000000000000000000000..a3f369847677d8dbaaadb8297691b1be92cf189f --- /dev/null +++ b/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() diff --git a/core/models.py b/core/models.py new file mode 100755 index 0000000000000000000000000000000000000000..fbc59ac019198da8c8f920794f57b3277cc6dfb3 --- /dev/null +++ b/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 +} diff --git a/core/utils/__pycache__/camera.cpython-310.pyc b/core/utils/__pycache__/camera.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..3acb23bb4626fe68b385d2e704f233a479da3ae6 Binary files /dev/null and b/core/utils/__pycache__/camera.cpython-310.pyc differ diff --git a/core/utils/__pycache__/dinov2.cpython-310.pyc b/core/utils/__pycache__/dinov2.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..c8296e1e783be73cfd273cd4481b6710a6084430 Binary files /dev/null and b/core/utils/__pycache__/dinov2.cpython-310.pyc differ diff --git a/core/utils/__pycache__/io.cpython-310.pyc b/core/utils/__pycache__/io.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..b5eca43ffa57537f82f7759c89a8ac69794f39f5 Binary files /dev/null and b/core/utils/__pycache__/io.cpython-310.pyc differ diff --git a/core/utils/__pycache__/math_utils.cpython-310.pyc b/core/utils/__pycache__/math_utils.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..e7086ba901c2251f670a388679944e23b2f26993 Binary files /dev/null and b/core/utils/__pycache__/math_utils.cpython-310.pyc differ diff --git a/core/utils/__pycache__/train_utils.cpython-310.pyc b/core/utils/__pycache__/train_utils.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..385718869cd7480895756de358e4b524fa76b279 Binary files /dev/null and b/core/utils/__pycache__/train_utils.cpython-310.pyc differ diff --git a/core/utils/__pycache__/vis_utils.cpython-310.pyc b/core/utils/__pycache__/vis_utils.cpython-310.pyc new file mode 100755 index 0000000000000000000000000000000000000000..b36ec14d8d14ce4a0444cb0d3e7ea7526805216c Binary files /dev/null and b/core/utils/__pycache__/vis_utils.cpython-310.pyc differ diff --git a/core/utils/camera.py b/core/utils/camera.py new file mode 100755 index 0000000000000000000000000000000000000000..a200b5ba440f2db7e84423152e3d241d6f55b529 --- /dev/null +++ b/core/utils/camera.py @@ -0,0 +1,53 @@ +import numpy as np +from scipy.spatial.transform import Rotation +import bpy + + +def setup_camera(cam_location, cam_rot=(0, 0, 0)): + bpy.ops.object.camera_add(location=cam_location, rotation=cam_rot) + camera = bpy.context.active_object + camera.rotation_mode = 'XYZ' + bpy.data.scenes["Scene"].camera = camera + scene = bpy.context.scene + camera.data.sensor_height = ( + camera.data.sensor_width * scene.render.resolution_y / scene.render.resolution_x + ) + for area in bpy.context.screen.areas: + if area.type == "VIEW_3D": + area.spaces.active.region_3d.view_perspective = "CAMERA" + break + cam_info_ng = bpy.data.node_groups.get("nodegroup_active_cam_info") + if cam_info_ng is not None: + cam_info_ng.nodes["Object Info"].inputs["Object"].default_value = camera + return camera + + +def convert_sphere_to_xyz(dist, elevation, azimuth): + """ + Convert spherical to cartesian coordinates. Assume camera is always looking at origin. + """ + elevation_rad = elevation * np.pi / 180.0 + azimuth_rad = azimuth * np.pi / 180.0 + cam_pos_x = dist * np.sin(elevation_rad) * np.cos(azimuth_rad) + cam_pos_y = dist * np.sin(elevation_rad) * np.sin(azimuth_rad) + cam_pos_z = dist * np.cos(elevation_rad) + # rotation + gaze_direction = -np.array([cam_pos_x, cam_pos_y, cam_pos_z]) + R_look_at = look_at_rotation_matrix(gaze_direction) + cam_rot_x, cam_rot_y, cam_rot_z = rotation_matrix_to_euler_angles(R_look_at) + return [cam_pos_x, cam_pos_y, cam_pos_z, cam_rot_x, cam_rot_y, cam_rot_z] + +def look_at_rotation_matrix(gaze_direction, world_up=(0,0,1)): + forward = gaze_direction / np.linalg.norm(gaze_direction) + right = np.cross(forward, world_up) + right /= np.linalg.norm(right) + up = np.cross(right, forward) + up /= np.linalg.norm(up) + R_look_at = np.array([right, up, -forward]).T + return R_look_at + +def rotation_matrix_to_euler_angles(R): + r = Rotation.from_matrix(R) + angles = r.as_euler("xyz", degrees=False) + + return angles[0], angles[1], angles[2] diff --git a/core/utils/dinov2.py b/core/utils/dinov2.py new file mode 100755 index 0000000000000000000000000000000000000000..47d22dc0609258e9eb7842a0fbb840b5dd5b7ff5 --- /dev/null +++ b/core/utils/dinov2.py @@ -0,0 +1,45 @@ +import os +import numpy as np +import torch +import torchvision.transforms as transforms +from PIL import Image + +class Dinov2Model(object): + def __init__(self, device='cuda'): + self.device = device + self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14') + # To load from local directory + # self.model = torch.hub.load('/path/to/your/local/dinov/repo', 'dinov2_vitb14', source='local', pretrained=False) + # self.model.load_state_dict(torch.load('/path/to/your/local/weights')) + self.model.to(device) + self.model.eval() + self.image_transform = transforms.Compose( + [ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) + self.grid_size = 224 // self.model.patch_size + + def encode_img(self, img_path, background=0): + image = Image.open(img_path).convert('RGB') + if background == 0: + mask = (np.array(image).sum(-1) <= 3) + img_arr = np.array(image) + img_arr[mask] = [255, 255, 255] + image = Image.fromarray(img_arr) + image = self.image_transform(image).unsqueeze(0).to(self.device) + with torch.no_grad(): + image_feat = self.model(image).float() + return image_feat + + def encode_batch_imgs(self, batch_imgs, global_feat=True): + with torch.no_grad(): + images = [self.image_transform(Image.fromarray(img)).to(self.device) for img in batch_imgs] + images = torch.stack(images, 0) + if global_feat: + image_feat = self.model(images).float() + else: + image_feat = self.model.get_intermediate_layers(images)[0].float() + return image_feat diff --git a/core/utils/io.py b/core/utils/io.py new file mode 100755 index 0000000000000000000000000000000000000000..fa2d6244d6f39c30173ebfa2be550d95db83af47 --- /dev/null +++ b/core/utils/io.py @@ -0,0 +1,52 @@ +import numpy as np +import bpy +from PIL import Image + +def read_list_from_txt(path): + with open(path, 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + return lines + +def save_points_as_ply(points, filename): + # Define the PLY header + # Save the points to a PLY file + with open(filename, "w") as f: + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write("element vertex " + str(len(points)) + "\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + f.write("end_header\n") + for point in points: + f.write("{} {} {}\n".format(point[0], point[1], point[2])) + +def make_gif(save_name, image_list): + frames = [Image.open(image) for image in image_list] + frame_one = frames[0] + frame_one.save(save_name, save_all=True, append_images=frames[1:], fps=15, loop=0, disposal=2, optimize=False, lossless=True) + +def clean_scene(): + bpy.ops.object.select_all(action='SELECT') + bpy_data = [bpy.data.actions, + bpy.data.armatures, + bpy.data.brushes, + bpy.data.cameras, + bpy.data.materials, + bpy.data.meshes, + bpy.data.objects, + bpy.data.shape_keys, + bpy.data.textures, + bpy.data.collections, + bpy.data.node_groups, + bpy.data.images, + bpy.data.movieclips, + bpy.data.curves, + bpy.data.particles] + + for bpy_data_iter in bpy_data: + for data in bpy_data_iter: + bpy_data_iter.remove(data, do_unlink=True) + + diff --git a/core/utils/math_utils.py b/core/utils/math_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..2946297d8112896cee38306b829f3d697be93623 --- /dev/null +++ b/core/utils/math_utils.py @@ -0,0 +1,71 @@ +import math +import numpy as np + +def norm_pdf(x, mean, sd): + var = float(sd)**2 + denom = (2 * math.pi * var)**.5 + num = math.exp(-(float(x) - float(mean))**2 / (2 * var)) + return num / denom + +def norm_cdf(x, mean, sd): + # calculate the cumulative distribution function for the normal distribution + return (1. + math.erf((x - mean) / (math.sqrt(2) * sd))) / 2. + +def normalize_params(params, params_dict): + # 1. normalize the GT params into [-1, 1] range 2. convert discrete params into continuous + keys_p, keys_d = params.keys(), params_dict.keys() + assert set(keys_p) == set(keys_d) + for key in params.keys(): + param_type = params_dict[key][0] + if param_type == "discrete": + # assume discrete params are represented as continuous integers + choices = params_dict[key][1] + leng = len(choices) + if choices[0].__class__ == float: + # TODO: this is to fix float discrete params + idx = np.where(np.array(choices) == float(params[key]))[0][0] + else: + idx = np.where(np.array(choices) == float(int(params[key])))[0][0] + # uniformly partition the target [-1, 1] range into equal parts + # set the discrete value as the middle point of the partition + params[key] = 2.0 / leng * (idx + 0.5) - 1 + elif param_type == "continuous": + # uniformly project the parameter value into [-1, 1] range + min_v, max_v = params_dict[key][1] + params[key] = ((params[key] - min_v) / (max_v - min_v)) * 2 - 1 + elif param_type == "normal": + mean, std = params_dict[key][1] + # clamp the normal distribution to [mean-3\sigma, mean+3\sigma], + # and then uniformly project the value into [-1, 1] range + min_v, max_v = mean - 3 * std, mean + 3 * std + params[key] = ((params[key] - min_v) / (max_v - min_v)) * 2 - 1 + else: + raise NotImplementedError + return params + +def unnormalize_params(params, params_dict): + # project the parameters back to the original range + # TODO: assume the params is ordered in the same order as params_dict keys + keys = params_dict.keys() + params_u = {} + for i, key in enumerate(keys): + param_type = params_dict[key][0] + # do the clamp to the predicted params into [-1, 1] + params_i = np.clip(params[i], -1, 1) + if param_type == "discrete": + choices = params_dict[key][1] + leng = len(choices) + idx = (params_i + 1) // (2 / leng) + if idx > leng - 1: + idx = leng - 1 + params_u[key] = choices[int(idx)] + elif param_type == "continuous": + min_v, max_v = params_dict[key][1] + params_u[key] = ((params_i + 1) / 2) * (max_v - min_v) + min_v + elif param_type == "normal": + mean, std = params_dict[key][1] + min_v, max_v = mean - 3 * std, mean + 3 * std + params_u[key] = ((params_i + 1) / 2) * (max_v - min_v) + min_v + else: + raise NotImplementedError + return params_u \ No newline at end of file diff --git a/core/utils/train_utils.py b/core/utils/train_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..9a5631a8882b91ceaf8e08cdae18822552afb52a --- /dev/null +++ b/core/utils/train_utils.py @@ -0,0 +1,70 @@ +import os +import torch +import numpy as np +import logging +from collections import OrderedDict +from PIL import Image + +def requires_grad(model, flag=True): + """ + Set requires_grad flag for all parameters in a model. + """ + for p in model.parameters(): + p.requires_grad = flag + +def create_logger(logging_dir): + """ + Create a logger that writes to a log file and stdout. + """ + logging.basicConfig( + level=logging.INFO, + format='[\033[34m%(asctime)s\033[0m] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] + ) + logger = logging.getLogger(__name__) + return logger + +@torch.no_grad() +def update_ema(ema_model, model, decay=0.9999): + """ + Step the EMA model towards the current model. + """ + ema_params = OrderedDict(ema_model.named_parameters()) + model_params = OrderedDict(model.named_parameters()) + + for name, param in model_params.items(): + name = name.replace("module.", "") + # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed + ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) + +def center_crop_arr(pil_image, image_size): + """ + Center cropping implementation from ADM. + https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 + """ + while min(*pil_image.size) >= 2 * image_size: + pil_image = pil_image.resize( + tuple(x // 2 for x in pil_image.size), resample=Image.BOX + ) + + scale = image_size / min(*pil_image.size) + pil_image = pil_image.resize( + tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC + ) + + arr = np.array(pil_image) + crop_y = (arr.shape[0] - image_size) // 2 + crop_x = (arr.shape[1] - image_size) // 2 + return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]) + +def load_model(ckpt_name): + """ + Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path. + """ + # Load a custom DiT checkpoint: + assert os.path.isfile(ckpt_name), f'Could not find DiT checkpoint at {ckpt_name}' + checkpoint = torch.load(ckpt_name, map_location=lambda storage, loc: storage) + if "ema" in checkpoint: # supports checkpoints from train.py + checkpoint = checkpoint["ema"] + return checkpoint \ No newline at end of file diff --git a/core/utils/vis_utils.py b/core/utils/vis_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..51af27b49fe570d5d35f3a65b66b80e49cdbfb56 --- /dev/null +++ b/core/utils/vis_utils.py @@ -0,0 +1,72 @@ +# Copyright 2020 Hsueh-Ti Derek Liu +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import bpy + +def setMat_plastic(mesh, meshColor, AOStrength = 0.0): + mat = bpy.data.materials.new('MeshMaterial') + mesh.data.materials.append(mat) + mesh.active_material = mat + mat.use_nodes = True + tree = mat.node_tree + + # set principled BSDF + tree.nodes["Principled BSDF"].inputs['Roughness'].default_value = 0.3 + #tree.nodes["Principled BSDF"].inputs['Sheen Tint'].default_value = [0, 0, 0, 1] + tree.nodes["Principled BSDF"].inputs['Sheen Tint'].default_value = 0 + tree.nodes["Principled BSDF"].inputs['Specular'].default_value = 0.5 + tree.nodes["Principled BSDF"].inputs['IOR'].default_value = 1.45 + tree.nodes["Principled BSDF"].inputs['Transmission'].default_value = 0 + tree.nodes["Principled BSDF"].inputs['Clearcoat Roughness'].default_value = 0 + + # add Ambient Occlusion + tree.nodes.new('ShaderNodeAmbientOcclusion') + tree.nodes.new('ShaderNodeGamma') + MIXRGB = tree.nodes.new('ShaderNodeMixRGB') + MIXRGB.blend_type = 'MULTIPLY' + tree.nodes["Gamma"].inputs["Gamma"].default_value = AOStrength + tree.nodes["Ambient Occlusion"].inputs["Distance"].default_value = 10.0 + tree.nodes["Gamma"].location.x -= 600 + + # set color using Hue/Saturation node + HSVNode = tree.nodes.new('ShaderNodeHueSaturation') + HSVNode.inputs['Color'].default_value = meshColor.RGBA + HSVNode.inputs['Saturation'].default_value = meshColor.S + HSVNode.inputs['Value'].default_value = meshColor.V + HSVNode.inputs['Hue'].default_value = meshColor.H + HSVNode.location.x -= 200 + + # set color brightness/contrast + BCNode = tree.nodes.new('ShaderNodeBrightContrast') + BCNode.inputs['Bright'].default_value = meshColor.B + BCNode.inputs['Contrast'].default_value = meshColor.C + BCNode.location.x -= 400 + + # link all the nodes + tree.links.new(HSVNode.outputs['Color'], BCNode.inputs['Color']) + tree.links.new(BCNode.outputs['Color'], tree.nodes['Ambient Occlusion'].inputs['Color']) + tree.links.new(tree.nodes["Ambient Occlusion"].outputs['Color'], MIXRGB.inputs['Color1']) + tree.links.new(tree.nodes["Ambient Occlusion"].outputs['AO'], tree.nodes['Gamma'].inputs['Color']) + tree.links.new(tree.nodes["Gamma"].outputs['Color'], MIXRGB.inputs['Color2']) + tree.links.new(MIXRGB.outputs['Color'], tree.nodes['Principled BSDF'].inputs['Base Color']) + +class colorObj(object): + def __init__(self, RGBA, \ + H = 0.5, S = 1.0, V = 1.0,\ + B = 0.0, C = 0.0): + self.H = H # hue + self.S = S # saturation + self.V = V # value + self.RGBA = RGBA + self.B = B # birghtness + self.C = C # contrast \ No newline at end of file diff --git a/examples/basket/001.png b/examples/basket/001.png new file mode 100755 index 0000000000000000000000000000000000000000..48a3acd5312816224039a835bd1d5caa4ef8d480 Binary files /dev/null and b/examples/basket/001.png differ diff --git a/examples/basket/002.png b/examples/basket/002.png new file mode 100755 index 0000000000000000000000000000000000000000..ec8ed0bb125fdf7f4f5b831a21a2e2825bfec0b1 Binary files /dev/null and b/examples/basket/002.png differ diff --git a/examples/basket/003.png b/examples/basket/003.png new file mode 100755 index 0000000000000000000000000000000000000000..b39e1befb7554813bb234a04ebf92336c3b289dc Binary files /dev/null and b/examples/basket/003.png differ diff --git a/examples/basket/004.png b/examples/basket/004.png new file mode 100755 index 0000000000000000000000000000000000000000..f145859fb6824643fdbba21b0a5081019a19f247 Binary files /dev/null and b/examples/basket/004.png differ diff --git a/examples/basket/005.png b/examples/basket/005.png new file mode 100755 index 0000000000000000000000000000000000000000..2c898cd741b83e95ff31d41ea9c37e5567ed397f Binary files /dev/null and b/examples/basket/005.png differ diff --git a/examples/chair/001.png b/examples/chair/001.png new file mode 100755 index 0000000000000000000000000000000000000000..6a477ebc59a715cd0b5d32a5acc785d9aefe2cc0 Binary files /dev/null and b/examples/chair/001.png differ diff --git a/examples/chair/002.png b/examples/chair/002.png new file mode 100755 index 0000000000000000000000000000000000000000..28eb51fc9b910278e678a77512b5678692e91cbd Binary files /dev/null and b/examples/chair/002.png differ diff --git a/examples/chair/003.png b/examples/chair/003.png new file mode 100755 index 0000000000000000000000000000000000000000..8a8984470690b3a3681eaa37f0e589e65e44a11a Binary files /dev/null and b/examples/chair/003.png differ diff --git a/examples/chair/004.png b/examples/chair/004.png new file mode 100755 index 0000000000000000000000000000000000000000..324ea82a7dbcfd97fcfa5ec37333ef70ca8eebab Binary files /dev/null and b/examples/chair/004.png differ diff --git a/examples/chair/005.png b/examples/chair/005.png new file mode 100755 index 0000000000000000000000000000000000000000..fac2274f66d3c5ccc0a150b39f9fc8c6762c9ca5 Binary files /dev/null and b/examples/chair/005.png differ diff --git a/examples/dandelion/001.png b/examples/dandelion/001.png new file mode 100755 index 0000000000000000000000000000000000000000..33efd85cb88f3149496cfbbc2634deddeb5aefe5 Binary files /dev/null and b/examples/dandelion/001.png differ diff --git a/examples/dandelion/002.png b/examples/dandelion/002.png new file mode 100755 index 0000000000000000000000000000000000000000..cbb78e077775edd1c9f70ad14019a85c056421be Binary files /dev/null and b/examples/dandelion/002.png differ diff --git a/examples/dandelion/003.png b/examples/dandelion/003.png new file mode 100755 index 0000000000000000000000000000000000000000..448954233fe08de672d0e6e468a69637d6a5b0e2 Binary files /dev/null and b/examples/dandelion/003.png differ diff --git a/examples/flower/001.png b/examples/flower/001.png new file mode 100755 index 0000000000000000000000000000000000000000..564a698c3b9d6ed8b8ee63f428c3e64dc6e9f552 Binary files /dev/null and b/examples/flower/001.png differ diff --git a/examples/flower/002.png b/examples/flower/002.png new file mode 100755 index 0000000000000000000000000000000000000000..e7a481c2af6825a0e68a736aaab7c6d56b0240e8 Binary files /dev/null and b/examples/flower/002.png differ diff --git a/examples/flower/003.png b/examples/flower/003.png new file mode 100755 index 0000000000000000000000000000000000000000..7cce7e65a57fc944e196862051e32488b4d01095 Binary files /dev/null and b/examples/flower/003.png differ diff --git a/examples/flower/004.png b/examples/flower/004.png new file mode 100755 index 0000000000000000000000000000000000000000..7a78eb7daa94716d5850aded7ddbdbace1f3ca2f Binary files /dev/null and b/examples/flower/004.png differ diff --git a/examples/table/001.png b/examples/table/001.png new file mode 100755 index 0000000000000000000000000000000000000000..ac5b76c4c64b60a28e2bbd2b4793c1c588b0ec8a Binary files /dev/null and b/examples/table/001.png differ diff --git a/examples/table/002.png b/examples/table/002.png new file mode 100755 index 0000000000000000000000000000000000000000..5f8b5eb3ad65a2d668f4ac4ae15ee979c7fdddd0 Binary files /dev/null and b/examples/table/002.png differ diff --git a/examples/table/003.png b/examples/table/003.png new file mode 100755 index 0000000000000000000000000000000000000000..a3d10e6e30e378eef769f5283817c813a130952e Binary files /dev/null and b/examples/table/003.png differ diff --git a/examples/table/004.png b/examples/table/004.png new file mode 100755 index 0000000000000000000000000000000000000000..9ef1bbf4deb2361eaeb0e23dbe0c8cc81b6c46fa Binary files /dev/null and b/examples/table/004.png differ diff --git a/examples/table/005.png b/examples/table/005.png new file mode 100755 index 0000000000000000000000000000000000000000..99c71b69af61b9c3d4bf581cbfa3a5a3378badcb Binary files /dev/null and b/examples/table/005.png differ diff --git a/examples/vase/001.png b/examples/vase/001.png new file mode 100755 index 0000000000000000000000000000000000000000..14edab60dd8c7b7c75e148e8da1b40c80feee888 Binary files /dev/null and b/examples/vase/001.png differ diff --git a/examples/vase/002.png b/examples/vase/002.png new file mode 100755 index 0000000000000000000000000000000000000000..d3277a0905e121c839dda196fd64e94ee10de048 Binary files /dev/null and b/examples/vase/002.png differ diff --git a/examples/vase/003.png b/examples/vase/003.png new file mode 100755 index 0000000000000000000000000000000000000000..d91e20b015ea77ce7b979b0312a26505f635fdb9 Binary files /dev/null and b/examples/vase/003.png differ diff --git a/examples/vase/004.png b/examples/vase/004.png new file mode 100755 index 0000000000000000000000000000000000000000..ad0b992c98e26775d5b252705e976763a9740453 Binary files /dev/null and b/examples/vase/004.png differ diff --git a/examples/vase/005.png b/examples/vase/005.png new file mode 100755 index 0000000000000000000000000000000000000000..831826eed3185ca6be4883b9ca9d3675af28df83 Binary files /dev/null and b/examples/vase/005.png differ diff --git a/package.txt b/package.txt new file mode 100755 index 0000000000000000000000000000000000000000..9e85fc6eaf7a81ec2462b2639e5ecfa52fc2a5c4 --- /dev/null +++ b/package.txt @@ -0,0 +1 @@ +libglm-dev \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100755 index 0000000000000000000000000000000000000000..9d84b62dcf75479a61f62254a45cae4a12494d56 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,38 @@ +rembg +onnxruntime-gpu==1.17 +torch==2.4.0 --index-url https://download.pytorch.org/whl/cu118 +torchvision==0.19.0 --index-url https://download.pytorch.org/whl/cu118 +torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu118 +accelerate==0.34.2 +bpy==3.6.0 +Cython==3.0.11 +diffusers==0.30.3 +einops==0.8.0 +gin-config==0.5.0 +imageio==2.35.1 +imgaug==0.4.0 +matplotlib==3.9.2 +matplotlib-inline==0.1.7 +omegaconf==2.3.0 +opencv-python==4.8.0.74 +opencv-python-headless==4.5.5.64 +pandas==2.2.2 +PyOpenGL==3.1.0 +pyrender==0.1.45 +PyYAML==6.0.2 +regex==2024.9.11 +safetensors==0.4.5 +scikit-image==0.24.0 +scikit-learn==1.5.1 +scipy==1.14.0 +six==1.16.0 +tensorboard==2.18.0 +tensorboard-data-server==0.7.2 +timm==0.4.12 +tokenizers==0.15.2 +tqdm==4.66.5 +transformers==4.36.0 +trimesh==4.4.4 +triton==3.0.0 +https://download.pytorch.org/whl/cu118/xformers-0.0.27.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl +./third_party/infinigen diff --git a/scripts/__pycache__/prepare_data.cpython-310.pyc b/scripts/__pycache__/prepare_data.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18c2541d014f80605b185d401ebd6803a0e4eaa1 Binary files /dev/null and b/scripts/__pycache__/prepare_data.cpython-310.pyc differ diff --git a/scripts/generate.py b/scripts/generate.py new file mode 100755 index 0000000000000000000000000000000000000000..84e993af51618a1ee1098dfb928c3f8f84ecc51b --- /dev/null +++ b/scripts/generate.py @@ -0,0 +1,152 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import numpy as np +import cv2 +import gin +import bpy +import gc +import logging +import argparse +from pathlib import Path +import importlib +import yaml + +logging.basicConfig( + format="[%(asctime)s.%(msecs)03d] [%(module)s] [%(levelname)s] | %(message)s", + datefmt="%H:%M:%S", + level=logging.INFO, +) +logger = logging.getLogger(__name__) + + +import infinigen +from infinigen.core import init, surface +from infinigen.assets.utils.decorate import read_co +from infinigen.assets.utils.misc import assign_material +from infinigen.core.util import blender as butil +from infinigen.assets.lighting import ( + hdri_lighting, + holdout_lighting, + sky_lighting, + three_point_lighting, +) +from core.utils.vis_utils import colorObj, setMat_plastic + + +def generate(generator, params, seed, mesh_save_path, no_mod=False, no_ground=True): + print(params) + print("Generating") + # reset to default + bpy.ops.wm.read_homefile(app_template="") + butil.clear_scene() + # Suppress info messages + bpy.ops.outliner.orphans_purge() + gc.collect() + # configurate infinigen + gin.add_config_file_search_path("./third_party/infinigen/infinigen_examples/configs_nature") + gin.parse_config_files_and_bindings( + ["configs/infinigen/base.gin"], + bindings=[], + skip_unknown=True, + finalize_config=False, + ) + surface.registry.initialize_from_gin() + print("Configured") + + # setup the scene + scene = bpy.context.scene + scene.render.engine = "CYCLES" + scene.cycles.device = "GPU" + scene.render.film_transparent = True + bpy.context.preferences.system.scrollback = 0 + bpy.context.preferences.edit.undo_steps = 0 + prefs = bpy.context.preferences.addons["cycles"].preferences + for dt in prefs.get_device_types(bpy.context): + prefs.get_devices_for_type(dt[0]) + bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" + + use_devices = [d for d in prefs.devices if d.type == "CUDA"] + for d in prefs.devices: + d.use = False + for d in use_devices: + d.use = True + + bpy.context.scene.world.node_tree.nodes["Background"].inputs[0].default_value[0:3] = (0.0, 0.0, 0.0) + print("Setup done") + + # update the parameters + generator.update_params(params) + # generate the object + asset = generator.spawn_asset(seed) + generator.finalize_assets(asset) + print("Generated") + + parent = asset + if asset.type == "EMPTY": + meshes = [o for o in asset.children_recursive if o.type == "MESH"] + sizes = [] + for m in meshes: + co = read_co(m) + sizes.append((np.amax(co, 0) - np.amin(co, 0)).sum()) + i = np.argmax(np.array(sizes)) + asset = meshes[i] + if not no_mod: + if parent.animation_data is not None: + drivers = parent.animation_data.drivers.values() + for d in drivers: + parent.driver_remove(d.data_path) + co = read_co(asset) + x_min, x_max = np.amin(co, 0), np.amax(co, 0) + parent.location = -(x_min[0] + x_max[0]) / 2, -(x_min[1] + x_max[1]) / 2, 0 + butil.apply_transform(parent, loc=True) + if not no_ground: + bpy.ops.mesh.primitive_grid_add( + size=5, x_subdivisions=400, y_subdivisions=400 + ) + plane = bpy.context.active_object + plane.location[-1] = x_min[-1] + plane.is_shadow_catcher = True + material = bpy.data.materials.new("plane") + material.use_nodes = True + material.node_tree.nodes["Principled BSDF"].inputs[0].default_value = ( + 0.015, + 0.009, + 0.003, + 1, + ) + assign_material(plane, material) + + # dump mesh model + save_mesh_filepath = mesh_save_path + bpy.ops.export_scene.gltf(filepath=save_mesh_filepath) + print("Mesh saved in {}".format(save_mesh_filepath)) + + +if __name__ == "__main__": + argparser = argparse.ArgumentParser() + argparser.add_argument("--config", type=str, required=True) + argparser.add_argument('--output_path', type=str, required=True) + argparser.add_argument('--params_path', type=str, required=True) + argparser.add_argument('--seed', type=int, default=0) + args = argparser.parse_args() + + # load config + with open(args.config) as f: + cfg = yaml.load(f, Loader=yaml.FullLoader) + + # load generators + generators_choices = ["chair", "table", "vase", "basket", "flower", "dandelion"] + factory_names = ["ChairFactory", "TableDiningFactory", "VaseFactory", "BasketBaseFactory", "FlowerFactory", "DandelionFactory"] + + category = args.config.split("/")[-1].split("_")[0] + idx = generators_choices.index(category) + + # load generator + module = importlib.import_module(f"core.assets.{category}") + gen = getattr(module, factory_names[idx]) + generator = gen(args.seed) + + # load params + params = np.load(args.params_path, allow_pickle=True).item() + generate(generator, params, args.seed, args.output_path) \ No newline at end of file diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py new file mode 100755 index 0000000000000000000000000000000000000000..7f96e317088bf9029307a39560311266a7ab3dfa --- /dev/null +++ b/scripts/prepare_data.py @@ -0,0 +1,472 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import numpy as np +import cv2 +import gin +import bpy +import gc +import logging +import time +import argparse +from pathlib import Path +import importlib +import json +import copy +import imgaug +import imgaug.augmenters as iaa +from core.utils.io import read_list_from_txt +from multiprocessing import Pool +import torch +from core.utils.dinov2 import Dinov2Model + +logging.basicConfig( + format="[%(asctime)s.%(msecs)03d] [%(module)s] [%(levelname)s] | %(message)s", + datefmt="%H:%M:%S", + level=logging.INFO, +) +logger = logging.getLogger(__name__) + + +import infinigen +from infinigen.core import init, surface +from infinigen.assets.utils.decorate import read_co +from infinigen.assets.utils.misc import assign_material +from infinigen.core.util import blender as butil +from infinigen.assets.lighting import ( + hdri_lighting, + holdout_lighting, + sky_lighting, + three_point_lighting, +) +from core.utils.camera import convert_sphere_to_xyz, setup_camera +from core.utils.vis_utils import colorObj, setMat_plastic + +# augment +color_aug = iaa.Sequential([ + # color aug + iaa.WithBrightnessChannels(iaa.Add((-50, 50))), + iaa.GammaContrast((0.7, 1.5), per_channel=True), + iaa.AddToHueAndSaturation((-60, 60)), + iaa.Grayscale((0.0, 0.8)), +]) +flip_aug = iaa.Sequential([iaa.Fliplr(0.5)]) +crop_aug = iaa.Sequential([ + iaa.CropAndPad(percent=(-0.1, 0.1), pad_mode='constant', pad_cval=(0, 0), keep_size=False), + iaa.CropToFixedSize(height=256, width=256), + iaa.PadToFixedSize(height=256, width=256) +]) +crop_resize_aug = iaa.KeepSizeByResize(iaa.Crop(percent=(0, 0.1), sample_independently=False, keep_size=False)) + +def aug(name, save_root, num_aug, flip=True, crop=True): + """Do the augmentation to RGBA image + """ + id, name = name.split("/") + img_rgba = cv2.imread(os.path.join(save_root, id, name), -1) # rgba + img = cv2.cvtColor(img_rgba[:,:,:3], cv2.COLOR_BGR2RGB) + img = np.concatenate([img, img_rgba[:,:,3:]], axis=2) # rgba + + save_dir = os.path.join(save_root, id) + os.makedirs(save_dir, exist_ok=True) + for j in range(num_aug): + # do the augmentation here + image, mask = copy.deepcopy(img[:,:,:3]), copy.deepcopy(img[:,:,3:]) + # aug color + if np.random.rand() < 0.8: + image = color_aug(image=image) + # flip + if flip: + image = flip_aug(image=np.concatenate([image, mask], axis=-1)) + else: + image = np.concatenate([image, mask], axis=-1) + # crop + if crop: + if np.random.rand() < 0.5: + # crop & pad + image = crop_aug(image=image) + else: + # crop & resize + image = crop_resize_aug(image=image) + if np.random.rand() < 0.1: + # binary image using mask + image, mask = image[:, :, :3], image[:, :, 3:] + # black image + image = np.tile(255 * (1.0 - (mask > 0)), (1,1,3)).astype(np.uint8) + image = np.concatenate([image, mask], axis=-1) + if np.random.rand() < 0.2: + image, mask = image[:, :, :3], image[:, :, 3:] + edge = np.expand_dims(cv2.Canny(mask, 100, 200), -1) + mask = (edge > 0).astype(np.uint8) + # convert edge into black + edge = 255 * (1.0 - mask) + image = np.tile(edge, (1,1,3)).astype(np.uint8) + image = np.concatenate([image, mask], axis=-1) + # save + save_name = os.path.join(save_dir, "{}_aug_{}.png".format(name[:-4], j)) + cv2.imwrite(save_name, image) + print(name) + + +def randomize_params(params_dict): + # Initialize the parameters + selected_params = {} + for key, value in params_dict.items(): + if value[0] == 'continuous': + min_v, max_v = value[1][0], value[1][1] + selected_params[key] = np.random.uniform(min_v, max_v) + elif value[0] == 'discrete': + choice_list = value[1] + selected_params[key] = np.random.choice(choice_list) + else: + raise NotImplementedError + return selected_params + +def generate(generator, params, seed, save_dir=None, save_name=None, + save_blend=False, save_img=False, save_untexture_img=False, save_gif=False, save_mesh=False, + cam_dists=[], cam_elevations=[], cam_azimuths=[], zoff=0, + resolution='256x256', sample=100, no_mod=False, no_ground=True, + window=None, screen=None): + print("Generating") + # reset to default + bpy.ops.wm.read_homefile(app_template="") + butil.clear_scene() + # Suppress info messages + bpy.ops.outliner.orphans_purge() + gc.collect() + # configurate infinigen + gin.add_config_file_search_path("./third_party/infinigen/infinigen_examples/configs_nature") + gin.parse_config_files_and_bindings( + ["configs/infinigen/base.gin"], + bindings=[], + skip_unknown=True, + finalize_config=False, + ) + surface.registry.initialize_from_gin() + + # setup the scene + scene = bpy.context.scene + scene.render.engine = "CYCLES" + scene.cycles.device = "GPU" + scene.render.film_transparent = True + bpy.context.preferences.system.scrollback = 0 + bpy.context.preferences.edit.undo_steps = 0 + prefs = bpy.context.preferences.addons["cycles"].preferences + for dt in prefs.get_device_types(bpy.context): + prefs.get_devices_for_type(dt[0]) + bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" + + use_devices = [d for d in prefs.devices if d.type == "CUDA"] + for d in prefs.devices: + d.use = False + for d in use_devices: + d.use = True + + scene.render.resolution_x, scene.render.resolution_y = map( + int, resolution.split("x") + ) + scene.cycles.samples = sample + bpy.context.scene.render.use_persistent_data = True + bpy.context.scene.world.node_tree.nodes["Background"].inputs[0].default_value[0:3] = (0.0, 0.0, 0.0) + + # update the parameters + generator.update_params(params) + # generate the object + asset = generator.spawn_asset(seed) + generator.finalize_assets(asset) + + parent = asset + if asset.type == "EMPTY": + meshes = [o for o in asset.children_recursive if o.type == "MESH"] + sizes = [] + for m in meshes: + co = read_co(m) + sizes.append((np.amax(co, 0) - np.amin(co, 0)).sum()) + i = np.argmax(np.array(sizes)) + asset = meshes[i] + if not no_mod: + if parent.animation_data is not None: + drivers = parent.animation_data.drivers.values() + for d in drivers: + parent.driver_remove(d.data_path) + co = read_co(asset) + x_min, x_max = np.amin(co, 0), np.amax(co, 0) + parent.location = -(x_min[0] + x_max[0]) / 2, -(x_min[1] + x_max[1]) / 2, 0 + butil.apply_transform(parent, loc=True) + if not no_ground: + bpy.ops.mesh.primitive_grid_add( + size=5, x_subdivisions=400, y_subdivisions=400 + ) + plane = bpy.context.active_object + plane.location[-1] = x_min[-1] + plane.is_shadow_catcher = True + material = bpy.data.materials.new("plane") + material.use_nodes = True + material.node_tree.nodes["Principled BSDF"].inputs[0].default_value = ( + 0.015, + 0.009, + 0.003, + 1, + ) + assign_material(plane, material) + + if save_blend: + # visualize the generated model by rendering + butil.save_blend(f"{save_dir}/{save_name}.blend", autopack=True) + + # render image + if save_img: + sky_lighting.add_lighting() + scene.render.image_settings.file_format = "PNG" + scene.render.image_settings.color_mode = "RGBA" + nodes = bpy.data.worlds["World"].node_tree.nodes + sky_texture = [n for n in nodes if n.name.startswith("Sky Texture")][-1] + sky_texture.sun_elevation = np.deg2rad(60) + sky_texture.sun_rotation = np.pi * 0.75 + + for cd in cam_dists: + for ce in cam_elevations: + for ca in cam_azimuths: + save_name_full = f"{save_name}_{cd}_{ce}_{ca}" + cam_data = convert_sphere_to_xyz(cd, ce, ca) + cam_location, cam_rot = cam_data[:3], cam_data[3:] + cam_location[-1] += zoff # TODO: fix the table case + camera = setup_camera(cam_location=cam_location, cam_rot=cam_rot) + cam_info_ng = bpy.data.node_groups.get("nodegroup_active_cam_info") + if cam_info_ng is not None: + cam_info_ng.nodes["Object Info"].inputs["Object"].default_value = camera + + image_path = str(f"{save_dir}/{save_name_full}_texture.png") + scene.render.filepath = image_path + bpy.ops.render.render(write_still=True) + # render untextured object + if save_untexture_img: + bpy.ops.object.shade_smooth() + asset.data.materials.clear() + # untextured model + #RGBA = (144.0/255, 210.0/255, 236.0/255, 1) + RGBA = (192.0/255, 192.0/255, 192.0/255, 1) + meshColor = colorObj(RGBA, 0.5, 1.0, 1.0, 0.0, 2.0) + setMat_plastic(asset, meshColor) + image_path = str(f"{save_dir}/{save_name_full}_geometry.png") + scene.render.filepath = image_path + bpy.ops.render.render(write_still=True) + + # render gif of object rotating + if save_gif: + save_gif_dir = os.path.join(save_dir, "gif") + os.makedirs(save_gif_dir, exist_ok=True) + bpy.context.scene.frame_end = 60 + asset_parent = asset if asset.parent is None else asset.parent + asset_parent.driver_add("rotation_euler")[-1].driver.expression = f"frame/{60 / (2 * np.pi * 1)}" + imgpath = str(f"{save_gif_dir}/{save_name}_###.png") + scene.render.filepath = str(imgpath) + bpy.ops.render.render(animation=True) + from core.utils.io import make_gif + all_imgpaths = [str(os.path.join(save_gif_dir, p)) for p in sorted(os.listdir(save_gif_dir)) if p.endswith('.png')] + make_gif(f"{save_gif_dir}/{save_name}.gif", all_imgpaths) + + # dump mesh model + if save_mesh: + save_mesh_filepath = os.path.join(save_dir, save_name+".glb") + bpy.ops.export_scene.gltf(filepath=save_mesh_filepath) + print("Mesh saved in {}".format(save_mesh_filepath)) + + return asset, image_path + +if __name__ == "__main__": + argparser = argparse.ArgumentParser() + argparser.add_argument('--generator', type=str, default='ChairFactory', + help='Supported generator: [ChairFactory, VaseFactory, TableDiningFactory, BasketBaseFactory, FlowerFactory, DandelionFactory]') + argparser.add_argument('--save_root', type=str, required=True) + argparser.add_argument('--total_num', type=int, default=20000) + argparser.add_argument('--aug_num', type=int, default=5) + argparser.add_argument('--batch_size', type=int, default=1000) + argparser.add_argument('--seed', type=int, default=0) + args = argparser.parse_args() + + # setup + """Training image rendering settings + Chair: cam_dists - [1.8, 2.0], elevations: [50, 60, 80], azimuths: [0, 30, 60, 80], zoff: 0.0 + Vase: cam_dists - [1.2, 1.6, 2.0], elevations: [60, 80, 90], azimuths: [0], zoff: 0.3 + Table: cam_dists - [5.0, 6.0], elevations: [60, 70], azimuths: [0, 30, 60], zoff: 0.1 + Flower: cam_dists - [3.0, 4.0], elevations: [20, 30, 50, 60], azimuths: [0], zoff: 0 + Dandelion: cam_dists - [3.0], elevations: [90], azimuths: [0], zoff: 0.5 + Basket: cam_dists - [1.2, 1.6], elevations: [50, 60, 70], azimuths: [30, 60], zoff: 0.0 + """ + np.random.seed(args.seed) + flip, crop = True, True + # Different training data rendering settings for different generators to improve efficiency + if args.generator == "ChairFactory": + cam_dists = [1.8, 2.0] + elevations = [50, 60, 80] + azimuths = [0, 30, 60, 80] + zoff = 0.0 + elif args.generator == "TableDiningFactory": + cam_dists = [5.0, 6.0] + elevations = [60, 70] + azimuths = [0, 30, 60, 90] + zoff = 0.1 + elif args.generator == "VaseFactory": + cam_dists = [1.2, 1.6, 2.0] + elevations = [60, 80, 90] + azimuths = [0] + zoff = 0.3 + elif args.generator == "BasketBaseFactory": + cam_dists = [1.2, 1.6] + elevations = [50, 60, 70] + azimuths = [30, 60] + zoff = 0.0 + elif args.generator == "FlowerFactory": + cam_dists = [2.0, 3.0, 4.0] + elevations = [30, 50, 60, 80] + azimuths = [0] + zoff = 0 + elif args.generator == "DandelionFactory": + cam_dists = [3.0] + elevations = [90] + azimuths = [0] + zoff = 0.5 + flip = False + + os.makedirs(args.save_root, exist_ok=True) + train_ratio = 0.9 + sample = 100 + resolution = '256x256' + + + # load the Blender procedural generator + OBJECTS_PATH = Path("./core/assets/") + assert OBJECTS_PATH.exists(), OBJECTS_PATH + generator = None + for subdir in sorted(list(OBJECTS_PATH.iterdir())): + clsname = subdir.name.split(".")[0].strip() + with gin.unlock_config(): + module = importlib.import_module(f"core.assets.{clsname}") + if hasattr(module, args.generator): + generator = getattr(module, args.generator) + logger.info(f"Found {args.generator} in {subdir}") + break + logger.debug(f"{args.generator} not found in {subdir}") + if generator is None: + raise ModuleNotFoundError(f"{args.generator} not Found.") + gen = generator(args.seed) + + # save params dict file + params_dict_file = f"{args.save_root}/params_dict.txt" + json.dump(gen.params_dict, open(params_dict_file, "w")) + + # generate data main loop + for i in range(args.total_num): + # sample parameters + params = randomize_params(gen.params_dict) + # fix dependent parameters + params_fix_unused = gen.fix_unused_params(params) + save_name = f"{i:05d}" + save_dir = f"{args.save_root}/{save_name}" + os.makedirs(save_dir, exist_ok=True) + # generate and save rendering - for training data, skip the blend file to save storage + if i < args.total_num * train_ratio: + save_blend = False + else: + save_blend = True + asset, img_path = generate(gen, params_fix_unused, args.seed, save_dir=save_dir, save_name=save_name, + save_blend=save_blend, save_img=True, cam_dists=cam_dists, + cam_elevations=elevations, cam_azimuths=azimuths, zoff=zoff, sample=sample, resolution=resolution) + # save the parameters + json.dump(params_fix_unused, open(f"{save_dir}/params.txt", "w"), default=str) + + if i % 100 == 0: + logger.info(f"{i} / {args.total_num} finished") + + # write filelist + f = open(os.path.join(args.save_root, "train_list_mv.txt"), "w") + total_num = args.total_num + for i in range(int(total_num * train_ratio)): + for cam_dist in cam_dists: + for elevation in elevations: + for azimuth in azimuths: + f.write( + "{:05d}/{:05d}_{}_{}_{}.png\n".format( + i, i, cam_dist, elevation, azimuth + ) + ) + f.close() + f = open(os.path.join(args.save_root, "test_list_mv.txt"), "w") + for i in range(int(total_num * train_ratio), total_num): + for cam_dist in cam_dists: + for elevation in elevations: + for azimuth in azimuths: + f.write( + "{:05d}/{:05d}_{}_{}_{}.png\n".format( + i, i, cam_dist, elevation, azimuth + ) + ) + f.close() + # do the augmentation + # main loop + image_list = read_list_from_txt(os.path.join(args.save_root, "train_list_mv.txt")) + print("Augmenting...Total data: {}".format(len(image_list))) + p = Pool(16) + for i, name in enumerate(image_list): + p.apply_async(aug, args=(name, args.save_root, args.aug_num, flip, crop)) + p.close() + p.join() + + # write the new list + f = open(os.path.join(args.save_root, "train_list_mv_withaug.txt"), "w") + for i in range(int(total_num * train_ratio)): + for cam_dist in cam_dists: + for elevation in elevations: + for azimuth in azimuths: + f.write( + "{:05d}/{:05d}_{}_{}_{}.png\n".format( + i, i, cam_dist, elevation, azimuth + ) + ) + for j in range(args.aug_num): + f.write( + "{:05d}/{:05d}_{}_{}_{}_aug_{}.png\n".format( + i, i, cam_dist, elevation, azimuth, j + ) + ) + f.close() + + # extract features + # Setup PyTorch: + torch.manual_seed(0) + torch.set_grad_enabled(False) + dinov2_model = Dinov2Model() + + # read image paths + with open(os.path.join(args.save_root, "train_list_mv_withaug.txt"), "r") as f: + image_paths = f.readlines() + with open(os.path.join(args.save_root, "test_list_mv.txt"), "r") as f: + test_image_paths = f.readlines() + image_paths = image_paths + test_image_paths + + image_paths = [os.path.join(args.save_root, path.strip()) for path in image_paths] + print(f"Number of images: {len(image_paths)}") + for i in range(0, len(image_paths), args.batch_size): + batch_paths = image_paths[i:i + args.batch_size] + # pre-process the image - RGBA to RGB with white background + batch_images = [] + for path in batch_paths: + image = cv2.imread(path, -1) + mask = (image[...,-1:] > 0) + image_rgb = cv2.cvtColor(image[...,:3], cv2.COLOR_BGR2RGB) + # resize if not 256 + if image.shape[0] != 256 or image.shape[1] != 256: + image_rgb = cv2.resize(image_rgb, (256, 256), interpolation=cv2.INTER_NEAREST) + mask = cv2.resize((255 * mask[:,:,0]).astype(np.uint8), (256, 256), interpolation=cv2.INTER_NEAREST) + mask = (mask > 128)[:,:,None] + # convert the transparent pixels to white background + image_whiteback = image_rgb * mask + 255 * (1 - mask) + batch_images.append(np.array(image_whiteback).astype(np.uint8)) + + + batch_features = dinov2_model.encode_batch_imgs(batch_images, global_feat=False).detach().cpu().numpy() + save_paths = [p.replace(".png", "_dino_token.npz") for p in batch_paths] + # save the features + for save_path, feature in zip(save_paths, batch_features): + np.savez_compressed(save_path, feature) + print(f"Extracted features for {i} images.") diff --git a/scripts/sample_diffusion.py b/scripts/sample_diffusion.py new file mode 100755 index 0000000000000000000000000000000000000000..4519cbfed0f0848bc65a0486e6e6096e35bc4785 --- /dev/null +++ b/scripts/sample_diffusion.py @@ -0,0 +1,137 @@ +# 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. + +""" +Sample new images from a pre-trained DiT. +""" +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import argparse +import yaml +import json +import numpy as np +from pathlib import Path +import gin +import importlib +import logging +import cv2 +from huggingface_hub import hf_hub_download + +logging.basicConfig( + format="[%(asctime)s.%(msecs)03d] [%(module)s] [%(levelname)s] | %(message)s", + datefmt="%H:%M:%S", + level=logging.INFO, +) +logger = logging.getLogger(__name__) + +import torch +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +from core.diffusion import create_diffusion +from core.models import DiT_models +from core.utils.train_utils import load_model +from core.utils.math_utils import unnormalize_params +from scripts.prepare_data import generate +from core.utils.dinov2 import Dinov2Model + +def main(cfg, generator): + # Setup PyTorch: + torch.manual_seed(cfg["seed"]) + torch.set_grad_enabled(False) + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Load model: + 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 = load_model(cfg["ckpt_path"]) + model.load_state_dict(state_dict) + model.eval() # important! + diffusion = create_diffusion(str(cfg["num_sampling_steps"])) + # feature model + feature_model = Dinov2Model() + + img_names = sorted(os.listdir(cfg["condition_img_dir"])) + for name in img_names: + img_path = os.path.join(cfg["condition_img_dir"], name) + # Load condition image and extract features + img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) + # pre-process: resize to 256x256 + img = cv2.resize(img, (256, 256)) + img = np.array(img).astype(np.uint8) + + img_feat = feature_model.encode_batch_imgs([img], global_feat=False) + if len(img_feat.shape) == 2: + img_feat = img_feat.unsqueeze(1) + + # Create sampling noise: + 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) + + # save params + json.dump(params_original, open("{}/{}_params.txt".format(cfg["save_dir"], name), "w"), default=str) + + # generate 3D using sampled params + asset, _ = generate(generator, params_original, seed=cfg["seed"], save_dir=cfg["save_dir"], save_name=name, + save_blend=True, save_img=True, save_untexture_img=True, save_gif=False, save_mesh=True, + cam_dists=cfg["r_cam_dists"], cam_elevations=cfg["r_cam_elevations"], cam_azimuths=cfg["r_cam_azimuths"], zoff=cfg["r_zoff"], + resolution='720x720', sample=200) + print("Generating model using sampled parameters. Saved in {}".format(cfg["save_dir"])) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--remove_bg", type=bool, default=False) + args = parser.parse_args() + with open(args.config) as f: + cfg = yaml.load(f, Loader=yaml.FullLoader) + cfg["remove_bg"] = args.remove_bg + + # load the Blender procedural generator + OBJECTS_PATH = Path(cfg["generator_root"]) + assert OBJECTS_PATH.exists(), OBJECTS_PATH + generator = None + for subdir in sorted(list(OBJECTS_PATH.iterdir())): + clsname = subdir.name.split(".")[0].strip() + with gin.unlock_config(): + module = importlib.import_module(f"core.assets.{clsname}") + if hasattr(module, cfg["generator"]): + generator = getattr(module, cfg["generator"]) + logger.info("Found {} in {}".format(cfg["generator"], subdir)) + break + logger.debug("{} not found in {}".format(cfg["generator"], subdir)) + if generator is None: + raise ModuleNotFoundError("{} not Found.".format(cfg["generator"])) + gen = generator(cfg["seed"]) + # create visualize dir + os.makedirs(cfg["save_dir"], exist_ok=True) + main(cfg, gen) diff --git a/scripts/test_diffusion.py b/scripts/test_diffusion.py new file mode 100755 index 0000000000000000000000000000000000000000..ea0777f2c4c0e48168d348d188a99d7e7ea6d5d3 --- /dev/null +++ b/scripts/test_diffusion.py @@ -0,0 +1,155 @@ +# 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. + +""" +Sample new images from a pre-trained DiT. +""" +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import argparse +import yaml +import json +import numpy as np +from pathlib import Path +import gin +import importlib +import logging +import cv2 +import matplotlib.pyplot as plt + + +logging.basicConfig( + format="[%(asctime)s.%(msecs)03d] [%(module)s] [%(levelname)s] | %(message)s", + datefmt="%H:%M:%S", + level=logging.INFO, +) +logger = logging.getLogger(__name__) + +import torch +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +from torch.utils.data import DataLoader + +from core.diffusion import create_diffusion +from core.models import DiT_models +from core.dataset import ImageParamsDataset +from core.utils.train_utils import load_model +from core.utils.math_utils import unnormalize_params +from scripts.prepare_data import generate + +def main(cfg, generator): + # Setup PyTorch: + torch.manual_seed(cfg["seed"]) + torch.set_grad_enabled(False) + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Load model: + latent_size = cfg["num_params"] + model = DiT_models[cfg["model"]](input_size=latent_size).to(device) + # load a custom DiT checkpoint from train.py: + state_dict = load_model(cfg["ckpt_path"]) + model.load_state_dict(state_dict) + model.eval() # important! + diffusion = create_diffusion(str(cfg["num_sampling_steps"])) + + # Load dataset + dataset = ImageParamsDataset(cfg["data_root"], cfg["test_file"], cfg["params_dict_file"]) + loader = DataLoader( + dataset, + batch_size=cfg["batch_size"], + shuffle=False, + num_workers=cfg["num_workers"], + pin_memory=True, + drop_last=False + ) + params_dict = json.load(open(cfg["params_dict_file"])) + idx = 0 + total_error = np.zeros(cfg["num_params"]) + for x, img_feat, img in loader: + # sample from random noise, conditioned on image features + img_feat = img_feat.to(device) + + model_kwargs = dict(y=img_feat) + + z = torch.randn(cfg["batch_size"], 1, latent_size, device=device) + + # 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.reshape(cfg["batch_size"], 1, -1) + samples = samples.squeeze(1).cpu().numpy() + x = x.squeeze(1).cpu().numpy() + img = img.cpu().numpy() + if cfg["run_generate"]: + # save GT & sampled params & images + for x_, params, img_ in zip(x, samples, img): + # generate 3D using sampled params + params_original = unnormalize_params(params, params_dict) + save_dir = os.path.join(cfg["save_dir"], "{:05d}".format(idx)) + os.makedirs(save_dir, exist_ok=True) + save_name = "sampled" + asset, _ = generate(generator, params_original, seed=cfg["seed"], save_dir=save_dir, save_name=save_name, + save_blend=True, save_img=True, save_gif=False, save_mesh=True, + cam_dists=cfg["r_cam_dists"], cam_elevations=cfg["r_cam_elevations"], cam_azimuths=cfg["r_cam_azimuths"], zoff=cfg["r_zoff"], + resolution='256x256', sample=100) + np.save(os.path.join(save_dir, "params.npy"), params_original) + print("Generating model using sampled parameters. Saved in {}".format(save_dir)) + # also save GT image & GT params + x_original = unnormalize_params(x_, params_dict) + np.save(os.path.join(save_dir, "gt_params.npy"), x_original) + cv2.imwrite(os.path.join(save_dir, "gt.png"), img_[:,:,::-1]) + idx += 1 + + # calculate metrics for sampled params & GT params + error = np.abs(x - samples) + total_error += error + + # print the average error for each parameter + avg_error = total_error / len(dataset) + param_names = params_dict.keys() + for param_name, error in zip(param_names, avg_error): + print(f"{param_name}: {error:.4f}") + # plot the error for each parameter + fig, ax = plt.subplots() + fig.set_size_inches(20, 15) + ax.barh(param_names, avg_error) + ax.set_xlabel("Average Error") + ax.set_ylabel("Parameters") + ax.set_title("Average Error for Each Parameter") + plt.yticks(fontsize=10) + fig.tight_layout() + fig.savefig(os.path.join(cfg["save_dir"], "avg_error.png")) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, required=True) + args = parser.parse_args() + with open(args.config) as f: + cfg = yaml.load(f, Loader=yaml.FullLoader) + + # load the Blender procedural generator + OBJECTS_PATH = Path(cfg["generator_root"]) + assert OBJECTS_PATH.exists(), OBJECTS_PATH + generator = None + for subdir in sorted(list(OBJECTS_PATH.iterdir())): + clsname = subdir.name.split(".")[0].strip() + with gin.unlock_config(): + module = importlib.import_module(f"core.assets.{clsname}") + if hasattr(module, cfg["generator"]): + generator = getattr(module, cfg["generator"]) + logger.info("Found {} in {}".format(cfg["generator"], subdir)) + break + logger.debug("{} not found in {}".format(cfg["generator"], subdir)) + if generator is None: + raise ModuleNotFoundError("{} not Found.".format(cfg["generator"])) + gen = generator(cfg["seed"]) + # create visualize dir + os.makedirs(cfg["save_dir"], exist_ok=True) + main(cfg, gen) diff --git a/scripts/train_diffusion.py b/scripts/train_diffusion.py new file mode 100755 index 0000000000000000000000000000000000000000..fafbb79eaf8dbe6b8d1eda30f49fc3767cd67bff --- /dev/null +++ b/scripts/train_diffusion.py @@ -0,0 +1,168 @@ +# 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. + +""" +A minimal training script for DiT. +""" +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import torch +# the first flag below was False when we tested this script but True makes A100 training a lot faster: +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +from torch.utils.data import DataLoader +import numpy as np +from copy import deepcopy +from glob import glob +from time import time +import argparse +import os +import yaml +from accelerate import Accelerator +from torch.utils.tensorboard import SummaryWriter + +from core.models import DiT_models +from core.diffusion import create_diffusion +from core.dataset import ImageParamsDataset +from core.utils.train_utils import create_logger, update_ema, requires_grad + +################################################################################# +# Training Loop # +################################################################################# + +def main(cfg): + """ + Trains a new DiT model. + """ + assert torch.cuda.is_available(), "Training currently requires at least one GPU." + + # Setup accelerator: + accelerator = Accelerator() + device = accelerator.device + + # Setup an experiment folder: + if accelerator.is_main_process: + os.makedirs(cfg["save_dir"], exist_ok=True) # Make results folder (holds all experiment subfolders) + save_dir = cfg["save_dir"] + experiment_index = len(glob(f"{save_dir}/*")) + experiment_dir = "{}/{:03d}-{}-{}-{}".format(save_dir, experiment_index, cfg["model"], cfg["epochs"], cfg["batch_size"]) # Create an experiment folder + checkpoint_dir = "{}/checkpoints".format(experiment_dir) # Stores saved model checkpoints + os.makedirs(checkpoint_dir, exist_ok=True) + logger = create_logger(experiment_dir) + logger.info(f"Experiment directory created at {experiment_dir}") + writer = SummaryWriter(experiment_dir) + + # Create model: + latent_size = cfg["num_params"] + condition_channels = 768 + model = DiT_models[cfg["model"]](input_size=latent_size, condition_channels=condition_channels) + # Note that parameter initialization is done within the DiT constructor + model = model.to(device) + ema = deepcopy(model).to(device) # Create an EMA of the model for use after training + requires_grad(ema, False) + diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule + if accelerator.is_main_process: + logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}") + + + # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper): + optimizer = torch.optim.AdamW(model.parameters(), lr=float(cfg["lr"]), weight_decay=0) + + # Setup data: + dataset = ImageParamsDataset(cfg["data_root"], cfg["train_file"], cfg["params_dict_file"]) + loader = DataLoader( + dataset, + batch_size=int(cfg["batch_size"] // accelerator.num_processes), + shuffle=True, + num_workers=cfg["num_workers"], + pin_memory=True, + drop_last=True + ) + if accelerator.is_main_process: + logger.info(f"Dataset contains {len(dataset):,} images") + + # Prepare models for training: + update_ema(ema, model, decay=0) # Ensure EMA is initialized with synced weights + model.train() # important! This enables embedding dropout for classifier-free guidance + ema.eval() # EMA model should always be in eval mode + model, optimizer, loader = accelerator.prepare(model, optimizer, loader) + + # Variables for monitoring/logging purposes: + train_steps = 0 + log_steps = 0 + running_loss = 0 + start_time = time() + + if accelerator.is_main_process: + logger.info("Training for {} epochs...".format(cfg["epochs"])) + + # main training loop + for epoch in range(int(cfg["epochs"])): + if accelerator.is_main_process: + logger.info(f"Beginning epoch {epoch}...") + for x, img_feat, img in loader: + # prepare the inputs + x = x.to(device) + img_feat = img_feat.to(device) + x = x.unsqueeze(dim=1) # [B, 1, N] + t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) + model_kwargs = dict(y=img_feat) + loss_dict = diffusion.training_losses(model, x, t, model_kwargs) + loss = loss_dict["loss"].mean() + optimizer.zero_grad() + accelerator.backward(loss) + optimizer.step() + update_ema(ema, model) + writer.add_scalar("train/loss", loss.item(), train_steps) + + # Log loss values: + running_loss += loss.item() + log_steps += 1 + train_steps += 1 + if train_steps % cfg["logging_iter"] == 0: + # Measure training speed: + torch.cuda.synchronize() + end_time = time() + steps_per_sec = log_steps / (end_time - start_time) + # Reduce loss history over all processes: + avg_loss = torch.tensor(running_loss / log_steps, device=device) + avg_loss = avg_loss.item() / accelerator.num_processes + if accelerator.is_main_process: + logger.info(f"(Step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}") + # Reset monitoring variables: + running_loss = 0 + log_steps = 0 + start_time = time() + + # Save DiT checkpoint: + if train_steps % cfg["ckpt_iter"] == 0 and train_steps > 0: + if accelerator.is_main_process: + checkpoint = { + "model": model.state_dict(), + "ema": ema.state_dict(), + "optimizer": optimizer.state_dict(), + "config": cfg, + } + checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" + torch.save(checkpoint, checkpoint_path) + logger.info(f"Saved checkpoint to {checkpoint_path}") + + model.eval() # important! This disables randomized embedding dropout + # do any sampling/FID calculation/etc. with ema (or model) in eval mode ... + + if accelerator.is_main_process: + writer.flush() + logger.info("Done!") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, required=True) + args = parser.parse_args() + with open(args.config) as f: + cfg = yaml.load(f, Loader=yaml.FullLoader) + main(cfg) diff --git a/third_party/infinigen b/third_party/infinigen new file mode 160000 index 0000000000000000000000000000000000000000..8d079368c3546bc4b4002dd4a91f649c210583eb --- /dev/null +++ b/third_party/infinigen @@ -0,0 +1 @@ +Subproject commit 8d079368c3546bc4b4002dd4a91f649c210583eb