Merge branch 'MrForExample:dev' into dev

.github/workflows/publish.yml

@@ -7,15 +7,19 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'MrForExample' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           ## Add your own personal access token to your Github Repository secrets and reference it here.
-          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

Configs/MV_Adapter_configs/geometry-guidance/mvadapter_ig2mv_partialimg_sdxl.yaml

@@ -0,0 +1,123 @@
+name: ig2mv
+tag: "r768-ortho-nv6-ig2mv-partialimg-dcrowcol-sdxl"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: data/texture_ortho6view_easylight
+  scene_list: data/objaverse_list_6w.json
+  background_color: gray
+  image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["position", "normal"]
+  position_offset: 0.5
+  position_scale: 1.0
+
+  reference_root_dir: ["data/texture_rand_easylight_objaverse"]
+  reference_scene_list: ["data/objaverse_list_6w.json"]
+  reference_image_modality: color
+  reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
+  reference_mask_aug: true
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv", "to_out_ref"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv", "_ref"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/geometry-guidance/mvadapter_ig2mv_sdxl.yaml

@@ -0,0 +1,122 @@
+name: ig2mv
+tag: "r768-ortho-nv6-ig2mv-dcrowcol-sdxl"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: data/texture_ortho6view_easylight
+  scene_list: data/objaverse_list_6w.json
+  background_color: gray
+  image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["position", "normal"]
+  position_offset: 0.5
+  position_scale: 1.0
+
+  reference_root_dir: ["data/texture_rand_easylight_objaverse"]
+  reference_scene_list: ["data/objaverse_list_6w.json"]
+  reference_image_modality: color
+  reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv", "to_out_ref"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv", "_ref"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/geometry-guidance/mvadapter_tg2mv_sdxl.yaml

@@ -0,0 +1,117 @@
+name: tg2mv
+tag: "r768-ortho-nv6-tg2mv-dcrowcol-sdxl"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: data/texture_ortho10view_easylight_objaverse
+  scene_list: data/objaverse_list_6w.json
+  background_color: gray
+  image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["position", "normal"]
+  position_offset: 0.5
+  position_scale: 1.0
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_text_sdxl.MVAdapterTextSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/view-guidance/mvadapter_i2mv_sd21.yaml

@@ -0,0 +1,122 @@
+name: i2mv
+tag: "r512-ortho-nv6-ele0-sd21"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: "data/texture_ortho10view_easylight_objaverse"
+  scene_list: "data/objaverse_list_6w.json"
+  background_color: gray
+  image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["plucker"]
+  plucker_offset: 1.0
+  plucker_scale: 2.0
+
+  reference_root_dir: ["data/texture_rand_easylight_objaverse"]
+  reference_scene_list: ["data/objaverse_list_6w.json"]
+  reference_image_modality: color
+  reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 512
+  width: 512
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_image_sd.MVAdapterImageSDSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-2-1-base"
+  pretrained_vae_name_or_path: null
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv", "to_out_ref"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv", "_ref"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 5
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/view-guidance/mvadapter_i2mv_sdxl.yaml

@@ -0,0 +1,122 @@
+name: i2mv
+tag: "r768-ortho-nv6-ele0-sdxl"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: "data/texture_ortho10view_easylight_objaverse"
+  scene_list: "data/objaverse_list_6w.json"
+  background_color: gray
+  image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["plucker"]
+  plucker_offset: 1.0
+  plucker_scale: 2.0
+
+  reference_root_dir: ["data/texture_rand_easylight_objaverse"]
+  reference_scene_list: ["data/objaverse_list_6w.json"]
+  reference_image_modality: color
+  reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv", "to_out_ref"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv", "_ref"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/view-guidance/mvadapter_i2mv_sdxl_aug_quantity.yaml

@@ -0,0 +1,123 @@
+name: i2mv
+tag: "r768-ortho-nv6-ele0-sdxl-aug-quantity"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: "data/texture_ortho10view_easylight_objaverse"
+  scene_list: "data/objaverse_list_6w.json"
+  background_color: gray
+  image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
+  image_modality: color
+  num_views: 6
+  random_view_list: [[0, 1, 2, 3, 4, 5], [0, 2, 3, 4], [0, 2, 3], [0, 3]]
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["plucker"]
+  plucker_offset: 1.0
+  plucker_scale: 2.0
+
+  reference_root_dir: ["data/texture_rand_easylight_objaverse"]
+  reference_scene_list: ["data/objaverse_list_6w.json"]
+  reference_image_modality: color
+  reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv", "to_out_ref"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv", "_ref"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/view-guidance/mvadapter_t2mv_sd21.yaml

@@ -0,0 +1,117 @@
+name: t2mv
+tag: "r512-ortho-nv6-ele0-sd"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: "data/texture_ortho10view_easylight_objaverse"
+  scene_list: "data/objaverse_list_6w.json"
+  background_color: gray
+  image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["plucker"]
+  plucker_offset: 1.0
+  plucker_scale: 2.0
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 512
+  width: 512
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_text_sd.MVAdapterTextSDSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-2-1-base"
+  pretrained_vae_name_or_path: null
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Configs/MV_Adapter_configs/view-guidance/mvadapter_t2mv_sdxl.yaml

@@ -0,0 +1,117 @@
+name: t2mv
+tag: "r768-ortho-nv6-ele0-sdxl"
+exp_root_dir: "outputs"
+seed: 42
+
+data_cls: mvadapter.data.multiview.MultiviewDataModule
+data:
+  root_dir: "data/texture_ortho10view_easylight_objaverse"
+  scene_list: "data/objaverse_list_6w.json"
+  background_color: gray
+  image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
+  image_modality: color
+  num_views: 6
+
+  prompt_db_path: data/objaverse_short_captions.json
+  return_prompt: true
+
+  projection_type: ORTHO
+
+  source_image_modality: ["plucker"]
+  plucker_offset: 1.0
+  plucker_scale: 2.0
+
+  train_indices: [0, -8]
+  val_indices: [-8, null]
+  test_indices: [-8, null]
+
+  height: 768
+  width: 768
+
+  batch_size: 1
+  num_workers: 16
+
+system_cls: mvadapter.systems.mvadapter_text_sdxl.MVAdapterTextSDXLSystem
+system:
+  check_train_every_n_steps: 1000
+  cleanup_after_validation_step: true
+  cleanup_after_test_step: true
+
+  # Model / Adapter
+  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+  pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
+  pretrained_adapter_name_or_path: null
+  init_adapter_kwargs:
+    # Multi-view adapter
+    self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
+    # Condition encoder
+    cond_in_channels: 6
+    # For training
+    copy_attn_weights: true
+    zero_init_module_keys: ["to_out_mv"]
+
+  # Training
+  train_cond_encoder: true
+  trainable_modules: ["_mv"]
+  prompt_drop_prob: 0.1
+  image_drop_prob: 0.1
+  cond_drop_prob: 0.1
+
+  # Noise sampler
+  shift_noise: true
+  shift_noise_mode: interpolated
+  shift_noise_scale: 8
+
+  # Evaluation
+  eval_seed: 42
+  eval_num_inference_steps: 30
+  eval_guidance_scale: 3.0
+  eval_height: ${data.height}
+  eval_width: ${data.width}
+
+  # optimizer definition
+  # you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5e-5
+      betas: [0.9, 0.999]
+      weight_decay: 0.01
+    params:
+      cond_encoder:
+        lr: 5e-5
+      unet:
+        lr: 5e-5
+
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+      - name: LinearLR
+        interval: step
+        args:
+          start_factor: 1e-6
+          end_factor: 1.0
+          total_iters: 2000
+      - name: ConstantLR
+        interval: step
+        args:
+          factor: 1.0
+          total_iters: 9999999
+    milestones: [2000]
+
+trainer:
+  max_epochs: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  val_check_interval: 2000
+  enable_progress_bar: true
+  precision: bf16-mixed
+  gradient_clip_val: 1.0
+  strategy: ddp
+  accumulate_grad_batches: 1
+
+checkpoint:
+  save_last: true # whether to save at each validation time
+  save_top_k: -1
+  every_n_epochs: 9999 # do not save at all for debug purpose

Gen_3D_Modules/MV_Adapter/mvadapter/loaders/__init__.py

@@ -0,0 +1 @@
+from .custom_adapter import CustomAdapterMixin

Gen_3D_Modules/MV_Adapter/mvadapter/loaders/custom_adapter.py

@@ -0,0 +1,132 @@
+import os
+from typing import Dict, Optional, Union
+
+import safetensors
+import torch
+from diffusers.utils import _get_model_file, logging
+from safetensors import safe_open
+from huggingface_hub import snapshot_download
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class CustomAdapterMixin:
+    def init_custom_adapter(self, *args, **kwargs):
+        self._init_custom_adapter(*args, **kwargs)
+
+    def _init_custom_adapter(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def load_custom_adapter(
+        self,
+        pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+        weight_name: str,
+        subfolder: Optional[str] = None,
+        local_cache_dir: Optional[str] = None,
+        **kwargs,
+    ):
+        # Load the main state dict first.
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+
+        user_agent = {
+            "file_type": "attn_procs_weights",
+            "framework": "pytorch",
+        }
+
+        if not isinstance(pretrained_model_name_or_path_or_dict, dict):
+            # First try to load from local cache if provided
+            if local_cache_dir and os.path.exists(os.path.join(local_cache_dir, weight_name)):
+                model_file = os.path.join(local_cache_dir, weight_name)
+                logger.info(f"Loading adapter from local cache: {model_file}")
+            else:
+                try:
+                    model_file = _get_model_file(
+                        pretrained_model_name_or_path_or_dict,
+                        weights_name=weight_name,
+                        subfolder=subfolder,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        proxies=proxies,
+                        local_files_only=local_files_only,
+                        token=token,
+                        revision=revision,
+                        user_agent=user_agent,
+                    )
+                except (OSError, EnvironmentError) as e:
+                    # If file not found and we have a custom cache dir, try to download
+                    if local_cache_dir:
+                        logger.info(f"Adapter not found. Downloading {pretrained_model_name_or_path_or_dict} to {local_cache_dir}")
+                        try:
+                            # Create directory if it doesn't exist
+                            os.makedirs(local_cache_dir, exist_ok=True)
+                            
+                            # Download the repository to local cache
+                            snapshot_download(
+                                repo_id=pretrained_model_name_or_path_or_dict,
+                                local_dir=local_cache_dir,
+                                ignore_patterns=["*.yaml", "*.json", "*.py", ".png", ".jpg", ".gif"]
+                            )
+                            
+                            # Now try to load from the local cache
+                            model_file = os.path.join(local_cache_dir, weight_name)
+                            if not os.path.exists(model_file):
+                                raise EnvironmentError(f"File {weight_name} not found even after download to {local_cache_dir}")
+                            
+                        except Exception as download_error:
+                            logger.error(f"Failed to download adapter: {download_error}")
+                            raise e  # Re-raise original error
+                    else:
+                        raise e  # Re-raise original error
+                
+            if weight_name.endswith(".safetensors"):
+                state_dict = {}
+                with safe_open(model_file, framework="pt", device="cpu") as f:
+                    for key in f.keys():
+                        state_dict[key] = f.get_tensor(key)
+            else:
+                state_dict = torch.load(model_file, map_location="cpu")
+        else:
+            state_dict = pretrained_model_name_or_path_or_dict
+
+        self._load_custom_adapter(state_dict)
+
+    def _load_custom_adapter(self, state_dict):
+        raise NotImplementedError
+
+    def save_custom_adapter(
+        self,
+        save_directory: Union[str, os.PathLike],
+        weight_name: str,
+        safe_serialization: bool = False,
+        **kwargs,
+    ):
+        if os.path.isfile(save_directory):
+            logger.error(
+                f"Provided path ({save_directory}) should be a directory, not a file"
+            )
+            return
+
+        if safe_serialization:
+
+            def save_function(weights, filename):
+                return safetensors.torch.save_file(
+                    weights, filename, metadata={"format": "pt"}
+                )
+
+        else:
+            save_function = torch.save
+
+        # Save the model
+        state_dict = self._save_custom_adapter(**kwargs)
+        save_function(state_dict, os.path.join(save_directory, weight_name))
+        logger.info(
+            f"Custom adapter weights saved in {os.path.join(save_directory, weight_name)}"
+        )
+
+    def _save_custom_adapter(self):
+        raise NotImplementedError

Gen_3D_Modules/MV_Adapter/mvadapter/models/attention_processor.py

@@ -0,0 +1,743 @@
+import math
+from typing import Callable, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention
+from diffusers.models.unets import UNet2DConditionModel
+from diffusers.utils import deprecate, logging
+from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
+from einops import rearrange, repeat
+from torch import nn
+
+
+def default_set_attn_proc_func(
+    name: str,
+    hidden_size: int,
+    cross_attention_dim: Optional[int],
+    ori_attn_proc: object,
+) -> object:
+    return ori_attn_proc
+
+
+def set_unet_2d_condition_attn_processor(
+    unet: UNet2DConditionModel,
+    set_self_attn_proc_func: Callable = default_set_attn_proc_func,
+    set_cross_attn_proc_func: Callable = default_set_attn_proc_func,
+    set_custom_attn_proc_func: Callable = default_set_attn_proc_func,
+    set_self_attn_module_names: Optional[List[str]] = None,
+    set_cross_attn_module_names: Optional[List[str]] = None,
+    set_custom_attn_module_names: Optional[List[str]] = None,
+) -> None:
+    do_set_processor = lambda name, module_names: (
+        any([name.startswith(module_name) for module_name in module_names])
+        if module_names is not None
+        else True
+    )  # prefix match
+
+    attn_procs = {}
+    for name, attn_processor in unet.attn_processors.items():
+        # set attn_processor by default, if module_names is None
+        set_self_attn_processor = do_set_processor(name, set_self_attn_module_names)
+        set_cross_attn_processor = do_set_processor(name, set_cross_attn_module_names)
+        set_custom_attn_processor = do_set_processor(name, set_custom_attn_module_names)
+
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+
+        is_custom = "attn_mid_blocks" in name or "attn_post_blocks" in name
+        if is_custom:
+            attn_procs[name] = (
+                set_custom_attn_proc_func(name, hidden_size, None, attn_processor)
+                if set_custom_attn_processor
+                else attn_processor
+            )
+        else:
+            cross_attention_dim = (
+                None
+                if name.endswith("attn1.processor")
+                else unet.config.cross_attention_dim
+            )
+            if cross_attention_dim is None or "motion_modules" in name:
+                # self attention
+                attn_procs[name] = (
+                    set_self_attn_proc_func(
+                        name, hidden_size, cross_attention_dim, attn_processor
+                    )
+                    if set_self_attn_processor
+                    else attn_processor
+                )
+            else:
+                # cross attention
+                attn_procs[name] = (
+                    set_cross_attn_proc_func(
+                        name, hidden_size, cross_attention_dim, attn_processor
+                    )
+                    if set_cross_attn_processor
+                    else attn_processor
+                )
+
+    unet.set_attn_processor(attn_procs)
+
+
+class DecoupledMVRowSelfAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for Decoupled Row-wise Self-Attention and Image Cross-Attention for PyTorch 2.0.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        inner_dim: int,
+        num_views: int = 1,
+        name: Optional[str] = None,
+        use_mv: bool = True,
+        use_ref: bool = False,
+    ):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "DecoupledMVRowSelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+        super().__init__()
+
+        self.num_views = num_views
+        self.name = name  # NOTE: need for image cross-attention
+        self.use_mv = use_mv
+        self.use_ref = use_ref
+
+        if self.use_mv:
+            self.to_q_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_k_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_v_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_out_mv = nn.ModuleList(
+                [
+                    nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
+                    nn.Dropout(0.0),
+                ]
+            )
+
+        if self.use_ref:
+            self.to_q_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_k_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_v_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_out_ref = nn.ModuleList(
+                [
+                    nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
+                    nn.Dropout(0.0),
+                ]
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        mv_scale: float = 1.0,
+        ref_hidden_states: Optional[torch.FloatTensor] = None,
+        ref_scale: float = 1.0,
+        cache_hidden_states: Optional[List[torch.FloatTensor]] = None,
+        use_mv: bool = True,
+        use_ref: bool = True,
+        num_views: Optional[int] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        """
+        New args:
+            mv_scale (float): scale for multi-view self-attention.
+            ref_hidden_states (torch.FloatTensor): reference encoder hidden states for image cross-attention.
+            ref_scale (float): scale for image cross-attention.
+            cache_hidden_states (List[torch.FloatTensor]): cache hidden states from reference unet.
+
+        """
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        if num_views is not None:
+            self.num_views = num_views
+
+        # NEW: cache hidden states for reference unet
+        if cache_hidden_states is not None:
+            cache_hidden_states[self.name] = hidden_states.clone()
+
+        # NEW: whether to use multi-view attention and image cross-attention
+        use_mv = self.use_mv and use_mv
+        use_ref = self.use_ref and use_ref
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        # NEW: for decoupled multi-view attention
+        if use_mv:
+            query_mv = self.to_q_mv(hidden_states)
+
+        # NEW: for decoupled reference cross attention
+        if use_ref:
+            query_ref = self.to_q_ref(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        ####### Decoupled multi-view self-attention ########
+        if use_mv:
+            key_mv = self.to_k_mv(encoder_hidden_states)
+            value_mv = self.to_v_mv(encoder_hidden_states)
+
+            query_mv = query_mv.view(batch_size, -1, attn.heads, head_dim)
+            key_mv = key_mv.view(batch_size, -1, attn.heads, head_dim)
+            value_mv = value_mv.view(batch_size, -1, attn.heads, head_dim)
+
+            height = width = math.isqrt(sequence_length)
+
+            # row self-attention
+            query_mv = rearrange(
+                query_mv,
+                "(b nv) (ih iw) h c -> (b nv ih) iw h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            ).transpose(1, 2)
+            key_mv = rearrange(
+                key_mv,
+                "(b nv) (ih iw) h c -> b ih (nv iw) h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            )
+            key_mv = (
+                key_mv.repeat_interleave(self.num_views, dim=0)
+                .view(batch_size * height, -1, attn.heads, head_dim)
+                .transpose(1, 2)
+            )
+            value_mv = rearrange(
+                value_mv,
+                "(b nv) (ih iw) h c -> b ih (nv iw) h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            )
+            value_mv = (
+                value_mv.repeat_interleave(self.num_views, dim=0)
+                .view(batch_size * height, -1, attn.heads, head_dim)
+                .transpose(1, 2)
+            )
+
+            hidden_states_mv = F.scaled_dot_product_attention(
+                query_mv,
+                key_mv,
+                value_mv,
+                dropout_p=0.0,
+                is_causal=False,
+            )
+            hidden_states_mv = rearrange(
+                hidden_states_mv,
+                "(b nv ih) h iw c -> (b nv) (ih iw) (h c)",
+                nv=self.num_views,
+                ih=height,
+            )
+            hidden_states_mv = hidden_states_mv.to(query.dtype)
+
+            # linear proj
+            hidden_states_mv = self.to_out_mv[0](hidden_states_mv)
+            # dropout
+            hidden_states_mv = self.to_out_mv[1](hidden_states_mv)
+
+        if use_ref:
+            reference_hidden_states = ref_hidden_states[self.name]
+
+            key_ref = self.to_k_ref(reference_hidden_states)
+            value_ref = self.to_v_ref(reference_hidden_states)
+
+            query_ref = query_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
+                1, 2
+            )
+            key_ref = key_ref.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            value_ref = value_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
+                1, 2
+            )
+
+            hidden_states_ref = F.scaled_dot_product_attention(
+                query_ref, key_ref, value_ref, dropout_p=0.0, is_causal=False
+            )
+
+            hidden_states_ref = hidden_states_ref.transpose(1, 2).reshape(
+                batch_size, -1, attn.heads * head_dim
+            )
+            hidden_states_ref = hidden_states_ref.to(query.dtype)
+
+            # linear proj
+            hidden_states_ref = self.to_out_ref[0](hidden_states_ref)
+            # dropout
+            hidden_states_ref = self.to_out_ref[1](hidden_states_ref)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if use_mv:
+            hidden_states = hidden_states + hidden_states_mv * mv_scale
+
+        if use_ref:
+            hidden_states = hidden_states + hidden_states_ref * ref_scale
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+    def set_num_views(self, num_views: int) -> None:
+        self.num_views = num_views
+
+
+class DecoupledMVRowColSelfAttnProcessor2_0(torch.nn.Module):
+    r"""
+    Attention processor for Decoupled Row-wise Self-Attention and Image Cross-Attention for PyTorch 2.0.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        inner_dim: int,
+        num_views: int = 1,
+        name: Optional[str] = None,
+        use_mv: bool = True,
+        use_ref: bool = False,
+    ):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "DecoupledMVRowSelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+        super().__init__()
+
+        self.num_views = num_views
+        self.name = name  # NOTE: need for image cross-attention
+        self.use_mv = use_mv
+        self.use_ref = use_ref
+
+        if self.use_mv:
+            self.to_q_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_k_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_v_mv = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_out_mv = nn.ModuleList(
+                [
+                    nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
+                    nn.Dropout(0.0),
+                ]
+            )
+
+        if self.use_ref:
+            self.to_q_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_k_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_v_ref = nn.Linear(
+                in_features=query_dim, out_features=inner_dim, bias=False
+            )
+            self.to_out_ref = nn.ModuleList(
+                [
+                    nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
+                    nn.Dropout(0.0),
+                ]
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        mv_scale: float = 1.0,
+        ref_hidden_states: Optional[torch.FloatTensor] = None,
+        ref_scale: float = 1.0,
+        cache_hidden_states: Optional[List[torch.FloatTensor]] = None,
+        use_mv: bool = True,
+        use_ref: bool = True,
+        num_views: Optional[int] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        """
+        New args:
+            mv_scale (float): scale for multi-view self-attention.
+            ref_hidden_states (torch.FloatTensor): reference encoder hidden states for image cross-attention.
+            ref_scale (float): scale for image cross-attention.
+            cache_hidden_states (List[torch.FloatTensor]): cache hidden states from reference unet.
+
+        """
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        if num_views is not None:
+            self.num_views = num_views
+
+        # NEW: cache hidden states for reference unet
+        if cache_hidden_states is not None:
+            cache_hidden_states[self.name] = hidden_states.clone()
+
+        # NEW: whether to use multi-view attention and image cross-attention
+        use_mv = self.use_mv and use_mv
+        use_ref = self.use_ref and use_ref
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        # NEW: for decoupled multi-view attention
+        if use_mv:
+            query_mv = self.to_q_mv(hidden_states)
+
+        # NEW: for decoupled reference cross attention
+        if use_ref:
+            query_ref = self.to_q_ref(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        ####### Decoupled multi-view self-attention ########
+        if use_mv:
+            key_mv = self.to_k_mv(encoder_hidden_states)
+            value_mv = self.to_v_mv(encoder_hidden_states)
+
+            query_mv = query_mv.view(batch_size, -1, attn.heads, head_dim)
+            key_mv = key_mv.view(batch_size, -1, attn.heads, head_dim)
+            value_mv = value_mv.view(batch_size, -1, attn.heads, head_dim)
+
+            height = width = math.isqrt(sequence_length)
+
+            query_mv = rearrange(
+                query_mv,
+                "(b nv) (ih iw) h c -> b nv ih iw h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            )
+            key_mv = rearrange(
+                key_mv,
+                "(b nv) (ih iw) h c -> b nv ih iw h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            )
+            value_mv = rearrange(
+                value_mv,
+                "(b nv) (ih iw) h c -> b nv ih iw h c",
+                nv=self.num_views,
+                ih=height,
+                iw=width,
+            )
+
+            # row-wise attention for view 0123 (front, right, back, left)
+            query_mv_0123 = rearrange(
+                query_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
+            )
+            key_mv_0123 = rearrange(
+                key_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
+            )
+            value_mv_0123 = rearrange(
+                value_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
+            )
+            hidden_states_mv_0123 = F.scaled_dot_product_attention(
+                query_mv_0123,
+                key_mv_0123,
+                value_mv_0123,
+                dropout_p=0.0,
+                is_causal=False,
+            )
+            hidden_states_mv_0123 = rearrange(
+                hidden_states_mv_0123,
+                "(b ih) h (nv iw) c -> b nv (ih iw) (h c)",
+                ih=height,
+                iw=height,
+            )
+
+            # col-wise attention for view 0245 (front, back, top, bottom)
+            # flip first
+            query_mv_0245 = torch.cat(
+                [
+                    torch.flip(query_mv[:, [0]], [3]),  # horizontal flip
+                    query_mv[:, [2, 4, 5]],
+                ],
+                dim=1,
+            )
+            key_mv_0245 = torch.cat(
+                [
+                    torch.flip(key_mv[:, [0]], [3]),  # horizontal flip
+                    key_mv[:, [2, 4, 5]],
+                ],
+                dim=1,
+            )
+            value_mv_0245 = torch.cat(
+                [
+                    torch.flip(value_mv[:, [0]], [3]),  # horizontal flip
+                    value_mv[:, [2, 4, 5]],
+                ],
+                dim=1,
+            )
+            # attention
+            query_mv_0245 = rearrange(
+                query_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c"
+            )
+            key_mv_0245 = rearrange(key_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c")
+            value_mv_0245 = rearrange(
+                value_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c"
+            )
+            hidden_states_mv_0245 = F.scaled_dot_product_attention(
+                query_mv_0245,
+                key_mv_0245,
+                value_mv_0245,
+                dropout_p=0.0,
+                is_causal=False,
+            )
+            # flip back
+            hidden_states_mv_0245 = rearrange(
+                hidden_states_mv_0245,
+                "(b iw) h (nv ih) c -> b nv ih iw (h c)",
+                ih=height,
+                iw=height,
+            )
+            hidden_states_mv_0245 = torch.cat(
+                [
+                    torch.flip(hidden_states_mv_0245[:, [0]], [3]),  # horizontal flip
+                    hidden_states_mv_0245[:, [1, 2, 3]],
+                ],
+                dim=1,
+            )
+            hidden_states_mv_0245 = hidden_states_mv_0245.view(
+                hidden_states_mv_0245.shape[0],
+                hidden_states_mv_0245.shape[1],
+                -1,
+                hidden_states_mv_0245.shape[-1],
+            )
+
+            # combine row and col
+            hidden_states_mv = torch.stack(
+                [
+                    (hidden_states_mv_0123[:, 0] + hidden_states_mv_0245[:, 0]) / 2,
+                    hidden_states_mv_0123[:, 1],
+                    (hidden_states_mv_0123[:, 2] + hidden_states_mv_0245[:, 1]) / 2,
+                    hidden_states_mv_0123[:, 3],
+                    hidden_states_mv_0245[:, 2],
+                    hidden_states_mv_0245[:, 3],
+                ],
+                dim=1,
+            )
+
+            hidden_states_mv = hidden_states_mv.view(
+                -1, hidden_states_mv.shape[-2], hidden_states_mv.shape[-1]
+            )
+            hidden_states_mv = hidden_states_mv.to(query.dtype)
+
+            # linear proj
+            hidden_states_mv = self.to_out_mv[0](hidden_states_mv)
+            # dropout
+            hidden_states_mv = self.to_out_mv[1](hidden_states_mv)
+
+        if use_ref:
+            reference_hidden_states = ref_hidden_states[self.name]
+
+            key_ref = self.to_k_ref(reference_hidden_states)
+            value_ref = self.to_v_ref(reference_hidden_states)
+
+            query_ref = query_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
+                1, 2
+            )
+            key_ref = key_ref.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            value_ref = value_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
+                1, 2
+            )
+
+            hidden_states_ref = F.scaled_dot_product_attention(
+                query_ref, key_ref, value_ref, dropout_p=0.0, is_causal=False
+            )
+
+            hidden_states_ref = hidden_states_ref.transpose(1, 2).reshape(
+                batch_size, -1, attn.heads * head_dim
+            )
+            hidden_states_ref = hidden_states_ref.to(query.dtype)
+
+            # linear proj
+            hidden_states_ref = self.to_out_ref[0](hidden_states_ref)
+            # dropout
+            hidden_states_ref = self.to_out_ref[1](hidden_states_ref)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if use_mv:
+            hidden_states = hidden_states + hidden_states_mv * mv_scale
+
+        if use_ref:
+            hidden_states = hidden_states + hidden_states_ref * ref_scale
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+    def set_num_views(self, num_views: int) -> None:
+        self.num_views = num_views

Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_i2mv_sdxl.py

@@ -0,0 +1,962 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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
+#
+#     http://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 inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import PIL
+import torch
+import torch.nn as nn
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.models import (
+    AutoencoderKL,
+    ImageProjection,
+    T2IAdapter,
+    UNet2DConditionModel,
+)
+from diffusers.pipelines.stable_diffusion_xl.pipeline_output import (
+    StableDiffusionXLPipelineOutput,
+)
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+    rescale_noise_cfg,
+    retrieve_timesteps,
+)
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from einops import rearrange
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
+
+from ..loaders import CustomAdapterMixin
+from ..models.attention_processor import (
+    DecoupledMVRowSelfAttnProcessor2_0,
+    set_unet_2d_condition_attn_processor,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor,
+    generator: Optional[torch.Generator] = None,
+    sample_mode: str = "sample",
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class MVAdapterI2MVSDXLPipeline(StableDiffusionXLPipeline, CustomAdapterMixin):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+        force_zeros_for_empty_prompt: bool = True,
+        add_watermarker: Optional[bool] = None,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            unet=unet,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
+            add_watermarker=add_watermarker,
+        )
+
+        self.control_image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_convert_rgb=True,
+            do_normalize=False,
+        )
+
+    # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.prepare_latents
+    def prepare_image_latents(
+        self,
+        image,
+        timestep,
+        batch_size,
+        num_images_per_prompt,
+        dtype,
+        device,
+        generator=None,
+        add_noise=True,
+    ):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        latents_mean = latents_std = None
+        if (
+            hasattr(self.vae.config, "latents_mean")
+            and self.vae.config.latents_mean is not None
+        ):
+            latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1)
+        if (
+            hasattr(self.vae.config, "latents_std")
+            and self.vae.config.latents_std is not None
+        ):
+            latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1)
+
+        # Offload text encoder if `enable_model_cpu_offload` was enabled
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.text_encoder_2.to("cpu")
+            torch.cuda.empty_cache()
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            if self.vae.config.force_upcast:
+                image = image.float()
+                self.vae.to(dtype=torch.float32)
+
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            elif isinstance(generator, list):
+                if image.shape[0] < batch_size and batch_size % image.shape[0] == 0:
+                    image = torch.cat([image] * (batch_size // image.shape[0]), dim=0)
+                elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0:
+                    raise ValueError(
+                        f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} "
+                    )
+
+                init_latents = [
+                    retrieve_latents(
+                        self.vae.encode(image[i : i + 1]), generator=generator[i]
+                    )
+                    for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = retrieve_latents(
+                    self.vae.encode(image), generator=generator
+                )
+
+            if self.vae.config.force_upcast:
+                self.vae.to(dtype)
+
+            init_latents = init_latents.to(dtype)
+            if latents_mean is not None and latents_std is not None:
+                latents_mean = latents_mean.to(device=device, dtype=dtype)
+                latents_std = latents_std.to(device=device, dtype=dtype)
+                init_latents = (
+                    (init_latents - latents_mean)
+                    * self.vae.config.scaling_factor
+                    / latents_std
+                )
+            else:
+                init_latents = self.vae.config.scaling_factor * init_latents
+
+        if (
+            batch_size > init_latents.shape[0]
+            and batch_size % init_latents.shape[0] == 0
+        ):
+            # expand init_latents for batch_size
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat(
+                [init_latents] * additional_image_per_prompt, dim=0
+            )
+        elif (
+            batch_size > init_latents.shape[0]
+            and batch_size % init_latents.shape[0] != 0
+        ):
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        if add_noise:
+            shape = init_latents.shape
+            noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            # get latents
+            init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+
+        latents = init_latents
+
+        return latents
+
+    def prepare_control_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance=False,
+        num_empty_images=0,  # for concat in batch like ImageDream
+    ):
+        assert hasattr(
+            self, "control_image_processor"
+        ), "control_image_processor is not initialized"
+
+        image = self.control_image_processor.preprocess(
+            image, height=height, width=width
+        ).to(dtype=torch.float32)
+
+        if num_empty_images > 0:
+            image = torch.cat(
+                [image, torch.zeros_like(image[:num_empty_images])], dim=0
+            )
+
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt  # always 1 for control image
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        if do_classifier_free_guidance:
+            image = torch.cat([image] * 2)
+
+        return image
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        # NEW
+        mv_scale: float = 1.0,
+        # Camera or geometry condition
+        control_image: Optional[PipelineImageInput] = None,
+        control_conditioning_scale: Optional[float] = 1.0,
+        control_conditioning_factor: float = 1.0,
+        # Image condition
+        reference_image: Optional[PipelineImageInput] = None,
+        reference_conditioning_scale: Optional[float] = 1.0,
+        **kwargs,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+                Anything below 512 pixels won't work well for
+                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
+                and checkpoints that are not specifically fine-tuned on low resolutions.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+                Anything below 512 pixels won't work well for
+                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
+                and checkpoints that are not specifically fine-tuned on low resolutions.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            denoising_end (`float`, *optional*):
+                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+                completed before it is intentionally prematurely terminated. As a result, the returned sample will
+                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
+                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
+                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
+                contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.0):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+                `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
+                explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                For most cases, `target_size` should be set to the desired height and width of the generated image. If
+                not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
+                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a specific image resolution. Part of SDXL's
+                micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
+                micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a target image resolution. It should be as same
+                as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            negative_prompt_2,
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._denoising_end = denoising_end
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        if negative_original_size is not None and negative_target_size is not None:
+            negative_add_time_ids = self._get_add_time_ids(
+                negative_original_size,
+                negative_crops_coords_top_left,
+                negative_target_size,
+                dtype=prompt_embeds.dtype,
+                text_encoder_projection_dim=text_encoder_projection_dim,
+            )
+        else:
+            negative_add_time_ids = add_time_ids
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat(
+                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
+            )
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(
+            batch_size * num_images_per_prompt, 1
+        )
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # Preprocess reference image
+        reference_image = self.image_processor.preprocess(reference_image)
+        reference_latents = self.prepare_image_latents(
+            reference_image,
+            timesteps[:1].repeat(batch_size * num_images_per_prompt),  # no use
+            batch_size,
+            1,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            add_noise=False,
+        )
+
+        with torch.no_grad():
+            ref_timesteps = torch.zeros_like(timesteps[0])
+            ref_hidden_states = {}
+
+            self.unet(
+                reference_latents,
+                ref_timesteps,
+                encoder_hidden_states=prompt_embeds[-1:],
+                added_cond_kwargs={
+                    "text_embeds": add_text_embeds[-1:],
+                    "time_ids": add_time_ids[-1:],
+                },
+                cross_attention_kwargs={
+                    "cache_hidden_states": ref_hidden_states,
+                    "use_mv": False,
+                    "use_ref": False,
+                },
+                return_dict=False,
+            )
+            ref_hidden_states = {
+                k: v.repeat_interleave(num_images_per_prompt, dim=0)
+                for k, v in ref_hidden_states.items()
+            }
+        if self.do_classifier_free_guidance:
+            ref_hidden_states = {
+                k: torch.cat([torch.zeros_like(v), v], dim=0)
+                for k, v in ref_hidden_states.items()
+            }
+
+        cross_attention_kwargs = {
+            "mv_scale": mv_scale,
+            "ref_hidden_states": {k: v.clone() for k, v in ref_hidden_states.items()},
+            "ref_scale": reference_conditioning_scale,
+            "num_views": num_images_per_prompt,
+            **(self.cross_attention_kwargs or {}),
+        }
+
+        # Preprocess control image
+        control_image_feature = self.prepare_control_image(
+            image=control_image,
+            width=width,
+            height=height,
+            batch_size=batch_size * num_images_per_prompt,
+            num_images_per_prompt=1,  # NOTE: always 1 for control images
+            device=device,
+            dtype=latents.dtype,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+        )
+        control_image_feature = control_image_feature.to(
+            device=device, dtype=latents.dtype
+        )
+
+        adapter_state = self.cond_encoder(control_image_feature)
+        for i, state in enumerate(adapter_state):
+            adapter_state[i] = state * control_conditioning_scale
+
+        # 8. Denoising loop
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+
+        # 8.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and isinstance(self.denoising_end, float)
+            and self.denoising_end > 0
+            and self.denoising_end < 1
+        ):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(
+                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
+            )
+            timesteps = timesteps[:num_inference_steps]
+
+        # 9. Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2)
+                    if self.do_classifier_free_guidance
+                    else latents
+                )
+
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+                added_cond_kwargs = {
+                    "text_embeds": add_text_embeds,
+                    "time_ids": add_time_ids,
+                }
+                if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+                    added_cond_kwargs["image_embeds"] = image_embeds
+
+                if i < int(num_inference_steps * control_conditioning_factor):
+                    down_intrablock_additional_residuals = [
+                        state.clone() for state in adapter_state
+                    ]
+                else:
+                    down_intrablock_additional_residuals = None
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(
+                        noise_pred,
+                        noise_pred_text,
+                        guidance_rescale=self.guidance_rescale,
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop(
+                        "negative_prompt_embeds", negative_prompt_embeds
+                    )
+                    add_text_embeds = callback_outputs.pop(
+                        "add_text_embeds", add_text_embeds
+                    )
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+                    negative_add_time_ids = callback_outputs.pop(
+                        "negative_add_time_ids", negative_add_time_ids
+                    )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        if not output_type == "latent":
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            needs_upcasting = (
+                self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+            )
+
+            if needs_upcasting:
+                self.upcast_vae()
+                latents = latents.to(
+                    next(iter(self.vae.post_quant_conv.parameters())).dtype
+                )
+            elif latents.dtype != self.vae.dtype:
+                if torch.backends.mps.is_available():
+                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                    self.vae = self.vae.to(latents.dtype)
+
+            # unscale/denormalize the latents
+            # denormalize with the mean and std if available and not None
+            has_latents_mean = (
+                hasattr(self.vae.config, "latents_mean")
+                and self.vae.config.latents_mean is not None
+            )
+            has_latents_std = (
+                hasattr(self.vae.config, "latents_std")
+                and self.vae.config.latents_std is not None
+            )
+            if has_latents_mean and has_latents_std:
+                latents_mean = (
+                    torch.tensor(self.vae.config.latents_mean)
+                    .view(1, 4, 1, 1)
+                    .to(latents.device, latents.dtype)
+                )
+                latents_std = (
+                    torch.tensor(self.vae.config.latents_std)
+                    .view(1, 4, 1, 1)
+                    .to(latents.device, latents.dtype)
+                )
+                latents = (
+                    latents * latents_std / self.vae.config.scaling_factor
+                    + latents_mean
+                )
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+        else:
+            image = latents
+
+        if not output_type == "latent":
+            # apply watermark if available
+            if self.watermark is not None:
+                image = self.watermark.apply_watermark(image)
+
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusionXLPipelineOutput(images=image)
+
+    ### NEW: adapters ###
+    def _init_custom_adapter(
+        self,
+        # Multi-view adapter
+        num_views: int = 1,
+        self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
+        # Condition encoder
+        cond_in_channels: int = 6,
+        # For training
+        copy_attn_weights: bool = True,
+        zero_init_module_keys: List[str] = [],
+    ):
+        # Condition encoder
+        self.cond_encoder = T2IAdapter(
+            in_channels=cond_in_channels,
+            channels=(320, 640, 1280, 1280),
+            num_res_blocks=2,
+            downscale_factor=16,
+            adapter_type="full_adapter_xl",
+        )
+
+        # set custom attn processor for multi-view attention and image cross-attention
+        self.unet: UNet2DConditionModel
+        set_unet_2d_condition_attn_processor(
+            self.unet,
+            set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
+                query_dim=hs,
+                inner_dim=hs,
+                num_views=num_views,
+                name=name,
+                use_mv=True,
+                use_ref=True,
+            ),
+            set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
+                query_dim=hs,
+                inner_dim=hs,
+                num_views=num_views,
+                name=name,
+                use_mv=False,
+                use_ref=False,
+            ),
+        )
+
+        # copy decoupled attention weights from original unet
+        if copy_attn_weights:
+            state_dict = self.unet.state_dict()
+            for key in state_dict.keys():
+                if "_mv" in key:
+                    compatible_key = key.replace("_mv", "").replace("processor.", "")
+                elif "_ref" in key:
+                    compatible_key = key.replace("_ref", "").replace("processor.", "")
+                else:
+                    compatible_key = key
+
+                is_zero_init_key = any([k in key for k in zero_init_module_keys])
+                if is_zero_init_key:
+                    state_dict[key] = torch.zeros_like(state_dict[compatible_key])
+                else:
+                    state_dict[key] = state_dict[compatible_key].clone()
+            self.unet.load_state_dict(state_dict)
+
+    def _load_custom_adapter(self, state_dict):
+        self.unet.load_state_dict(state_dict, strict=False)
+        self.cond_encoder.load_state_dict(state_dict, strict=False)
+
+    def _save_custom_adapter(
+        self,
+        include_keys: Optional[List[str]] = None,
+        exclude_keys: Optional[List[str]] = None,
+    ):
+        def include_fn(k):
+            is_included = False
+
+            if include_keys is not None:
+                is_included = is_included or any([key in k for key in include_keys])
+            if exclude_keys is not None:
+                is_included = is_included and not any(
+                    [key in k for key in exclude_keys]
+                )
+
+            return is_included
+
+        state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
+        state_dict.update(self.cond_encoder.state_dict())
+
+        return state_dict

Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_t2mv_sdxl.py

@@ -0,0 +1,801 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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
+#
+#     http://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.
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.models import AutoencoderKL, T2IAdapter, UNet2DConditionModel
+from diffusers.pipelines.stable_diffusion_xl.pipeline_output import (
+    StableDiffusionXLPipelineOutput,
+)
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
+    StableDiffusionXLPipeline,
+    rescale_noise_cfg,
+    retrieve_timesteps,
+)
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import deprecate, logging
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
+
+from ..loaders import CustomAdapterMixin
+from ..models.attention_processor import (
+    DecoupledMVRowSelfAttnProcessor2_0,
+    set_unet_2d_condition_attn_processor,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MVAdapterT2MVSDXLPipeline(StableDiffusionXLPipeline, CustomAdapterMixin):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+        force_zeros_for_empty_prompt: bool = True,
+        add_watermarker: Optional[bool] = None,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            unet=unet,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
+            add_watermarker=add_watermarker,
+        )
+
+        self.control_image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_convert_rgb=True,
+            do_normalize=False,
+        )
+
+    def prepare_control_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance=False,
+    ):
+        assert hasattr(
+            self, "control_image_processor"
+        ), "control_image_processor is not initialized"
+
+        image = self.control_image_processor.preprocess(
+            image, height=height, width=width
+        ).to(dtype=torch.float32)
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt  # always 1 for control image
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        if do_classifier_free_guidance:
+            image = torch.cat([image] * 2)
+
+        return image
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        # NEW
+        mv_scale: float = 1.0,
+        # Camera or geometry condition
+        control_image: Optional[PipelineImageInput] = None,
+        control_conditioning_scale: Optional[float] = 1.0,
+        control_conditioning_factor: float = 1.0,
+        # Optional. controlnet
+        controlnet_image: Optional[PipelineImageInput] = None,
+        controlnet_conditioning_scale: Optional[float] = 1.0,
+        **kwargs,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+                Anything below 512 pixels won't work well for
+                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
+                and checkpoints that are not specifically fine-tuned on low resolutions.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+                Anything below 512 pixels won't work well for
+                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
+                and checkpoints that are not specifically fine-tuned on low resolutions.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            denoising_end (`float`, *optional*):
+                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+                completed before it is intentionally prematurely terminated. As a result, the returned sample will
+                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
+                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
+                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
+                contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.0):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+                `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
+                explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                For most cases, `target_size` should be set to the desired height and width of the generated image. If
+                not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
+                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a specific image resolution. Part of SDXL's
+                micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
+                micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                To negatively condition the generation process based on a target image resolution. It should be as same
+                as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
+                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            negative_prompt_2,
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._denoising_end = denoising_end
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        if negative_original_size is not None and negative_target_size is not None:
+            negative_add_time_ids = self._get_add_time_ids(
+                negative_original_size,
+                negative_crops_coords_top_left,
+                negative_target_size,
+                dtype=prompt_embeds.dtype,
+                text_encoder_projection_dim=text_encoder_projection_dim,
+            )
+        else:
+            negative_add_time_ids = add_time_ids
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat(
+                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
+            )
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(
+            batch_size * num_images_per_prompt, 1
+        )
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # Preprocess control image
+        control_image_feature = self.prepare_control_image(
+            image=control_image,
+            width=width,
+            height=height,
+            batch_size=batch_size * num_images_per_prompt,
+            num_images_per_prompt=1,  # NOTE: always 1 for control images
+            device=device,
+            dtype=latents.dtype,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+        )
+        control_image_feature = control_image_feature.to(
+            device=device, dtype=latents.dtype
+        )
+
+        adapter_state = self.cond_encoder(control_image_feature)
+        for i, state in enumerate(adapter_state):
+            adapter_state[i] = state * control_conditioning_scale
+
+        # Preprocess controlnet image if provided
+        do_controlnet = controlnet_image is not None and hasattr(self, "controlnet")
+        if do_controlnet:
+            controlnet_image = self.prepare_control_image(
+                image=controlnet_image,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=1,  # NOTE: always 1 for control images
+                device=device,
+                dtype=latents.dtype,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+            )
+            controlnet_image = controlnet_image.to(device=device, dtype=latents.dtype)
+
+        # 8. Denoising loop
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+
+        # 8.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and isinstance(self.denoising_end, float)
+            and self.denoising_end > 0
+            and self.denoising_end < 1
+        ):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(
+                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
+            )
+            timesteps = timesteps[:num_inference_steps]
+
+        # 9. Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (
+                    torch.cat([latents] * 2)
+                    if self.do_classifier_free_guidance
+                    else latents
+                )
+
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+                added_cond_kwargs = {
+                    "text_embeds": add_text_embeds,
+                    "time_ids": add_time_ids,
+                }
+                if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+                    added_cond_kwargs["image_embeds"] = image_embeds
+
+                if i < int(num_inference_steps * control_conditioning_factor):
+                    down_intrablock_additional_residuals = [
+                        state.clone() for state in adapter_state
+                    ]
+                else:
+                    down_intrablock_additional_residuals = None
+
+                unet_add_kwargs = {}
+
+                # Do controlnet if provided
+                if do_controlnet:
+                    down_block_res_samples, mid_block_res_sample = self.controlnet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=prompt_embeds,
+                        controlnet_cond=controlnet_image,
+                        conditioning_scale=controlnet_conditioning_scale,
+                        guess_mode=False,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )
+                    unet_add_kwargs.update(
+                        {
+                            "down_block_additional_residuals": down_block_res_samples,
+                            "mid_block_additional_residual": mid_block_res_sample,
+                        }
+                    )
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs={
+                        "mv_scale": mv_scale,
+                        "num_views": num_images_per_prompt,
+                        **(self.cross_attention_kwargs or {}),
+                    },
+                    down_intrablock_additional_residuals=down_intrablock_additional_residuals,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                    **unet_add_kwargs,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(
+                        noise_pred,
+                        noise_pred_text,
+                        guidance_rescale=self.guidance_rescale,
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop(
+                        "negative_prompt_embeds", negative_prompt_embeds
+                    )
+                    add_text_embeds = callback_outputs.pop(
+                        "add_text_embeds", add_text_embeds
+                    )
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+                    negative_add_time_ids = callback_outputs.pop(
+                        "negative_add_time_ids", negative_add_time_ids
+                    )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        if not output_type == "latent":
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            needs_upcasting = (
+                self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+            )
+
+            if needs_upcasting:
+                self.upcast_vae()
+                latents = latents.to(
+                    next(iter(self.vae.post_quant_conv.parameters())).dtype
+                )
+            elif latents.dtype != self.vae.dtype:
+                if torch.backends.mps.is_available():
+                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                    self.vae = self.vae.to(latents.dtype)
+
+            # unscale/denormalize the latents
+            # denormalize with the mean and std if available and not None
+            has_latents_mean = (
+                hasattr(self.vae.config, "latents_mean")
+                and self.vae.config.latents_mean is not None
+            )
+            has_latents_std = (
+                hasattr(self.vae.config, "latents_std")
+                and self.vae.config.latents_std is not None
+            )
+            if has_latents_mean and has_latents_std:
+                latents_mean = (
+                    torch.tensor(self.vae.config.latents_mean)
+                    .view(1, 4, 1, 1)
+                    .to(latents.device, latents.dtype)
+                )
+                latents_std = (
+                    torch.tensor(self.vae.config.latents_std)
+                    .view(1, 4, 1, 1)
+                    .to(latents.device, latents.dtype)
+                )
+                latents = (
+                    latents * latents_std / self.vae.config.scaling_factor
+                    + latents_mean
+                )
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+        else:
+            image = latents
+
+        if not output_type == "latent":
+            # apply watermark if available
+            if self.watermark is not None:
+                image = self.watermark.apply_watermark(image)
+
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusionXLPipelineOutput(images=image)
+
+    ### NEW: adapters ###
+    def _init_custom_adapter(
+        self,
+        # Multi-view adapter
+        num_views: int = 1,
+        self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
+        # Condition encoder
+        cond_in_channels: int = 6,
+        # For training
+        copy_attn_weights: bool = True,
+        zero_init_module_keys: List[str] = [],
+    ):
+        # Condition encoder
+        self.cond_encoder = T2IAdapter(
+            in_channels=cond_in_channels,
+            channels=(320, 640, 1280, 1280),
+            num_res_blocks=2,
+            downscale_factor=16,
+            adapter_type="full_adapter_xl",
+        )
+
+        # set custom attn processor for multi-view attention
+        self.unet: UNet2DConditionModel
+        set_unet_2d_condition_attn_processor(
+            self.unet,
+            set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
+                query_dim=hs,
+                inner_dim=hs,
+                num_views=num_views,
+                name=name,
+                use_mv=True,
+                use_ref=False,
+            ),
+            set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
+                query_dim=hs,
+                inner_dim=hs,
+                num_views=num_views,
+                name=name,
+                use_mv=False,
+                use_ref=False,
+            ),
+        )
+
+        # copy decoupled attention weights from original unet
+        if copy_attn_weights:
+            state_dict = self.unet.state_dict()
+            for key in state_dict.keys():
+                if "_mv" in key:
+                    compatible_key = key.replace("_mv", "").replace("processor.", "")
+                else:
+                    compatible_key = key
+
+                is_zero_init_key = any([k in key for k in zero_init_module_keys])
+                if is_zero_init_key:
+                    state_dict[key] = torch.zeros_like(state_dict[compatible_key])
+                else:
+                    state_dict[key] = state_dict[compatible_key].clone()
+            self.unet.load_state_dict(state_dict)
+
+    def _load_custom_adapter(self, state_dict):
+        self.unet.load_state_dict(state_dict, strict=False)
+        self.cond_encoder.load_state_dict(state_dict, strict=False)
+
+    def _save_custom_adapter(
+        self,
+        include_keys: Optional[List[str]] = None,
+        exclude_keys: Optional[List[str]] = None,
+    ):
+        def include_fn(k):
+            is_included = False
+
+            if include_keys is not None:
+                is_included = is_included or any([key in k for key in include_keys])
+            if exclude_keys is not None:
+                is_included = is_included and not any(
+                    [key in k for key in exclude_keys]
+                )
+
+            return is_included
+
+        state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
+        state_dict.update(self.cond_encoder.state_dict())
+
+        return state_dict

Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_texture.py

@@ -0,0 +1,452 @@
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from spandrel import ModelLoader
+
+from MV_Adapter.mvadapter.utils import image_to_tensor, make_image_grid, tensor_to_image
+from MV_Adapter.mvadapter.utils.mesh_utils import (
+    Camera,
+    CameraProjection,
+    NVDiffRastContextWrapper,
+    SmartPainter,
+    TexturedMesh,
+    get_camera,
+    get_orthogonal_camera,
+    load_mesh,
+    render,
+    replace_mesh_texture_and_save,
+)
+from MV_Adapter.mvadapter.utils.mesh_utils.mesh_process import process_raw
+
+# Import dilate function to remove texture seams
+def torch_image_dilate(img):
+    """
+    Remove thin seams on generated texture
+        img (torch): [H, W, C]
+    """
+    import cv2
+    img = np.asarray(img.cpu().numpy(), dtype=np.float32)
+    img = img * 255
+    img = img.clip(0, 255)
+    mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)
+    mask = (mask <= 3.0).astype(np.float32)
+    kernel = np.ones((3, 3), 'uint8')
+    dilate_img = cv2.dilate(img, kernel, iterations=1)
+    img = img * (1 - mask) + dilate_img * mask
+    img = (img.clip(0, 255) / 255).astype(np.float32)
+    return torch.from_numpy(img)
+
+
+def clear():
+    torch.cuda.empty_cache()
+
+
+@contextmanager
+def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
+    texture_ = mesh.texture
+    mesh.texture = texture
+    try:
+        yield
+    finally:
+        mesh.texture = texture_
+
+
+@dataclass
+class ModProcessConfig:
+    view_upscale: bool = False
+    view_upscale_factor: int = 2
+    inpaint_mode: str = "uv"  # in ["none", "uv", "view"]
+    view_inpaint_max_view_score_thresh: float = 0.02
+    view_inpaint_min_rounds: int = 4
+    view_inpaint_max_rounds: int = 8
+    view_inpaint_uv_padding_end: bool = True
+
+
+@dataclass
+class TexturePipelineOutput:
+    shaded_model_save_path: Optional[str] = None
+    pbr_model_save_path: Optional[str] = None
+
+
+class TexturePipeline:
+    def __init__(self, upscaler_ckpt_path: str, inpaint_ckpt_path: str, device: str):
+        self.device = device
+        self.ctx = NVDiffRastContextWrapper(device=self.device)
+        self.cam_proj = CameraProjection(
+            pb_backend="torch-cuda", bg_remover=None, device=self.device
+        )
+        if upscaler_ckpt_path is not None:
+            self.upscaler = ModelLoader().load_from_file(upscaler_ckpt_path)
+            self.upscaler.to(self.device).eval().half()
+        if inpaint_ckpt_path is not None:
+            self.inpainter = ModelLoader().load_from_file(inpaint_ckpt_path)
+            self.inpainter.to(self.device).eval()
+
+        self.smart_painter = SmartPainter(self.device)
+
+    def load_packed_images(self, packed_image_path: Optional[str]) -> List[Image.Image]:
+        if packed_image_path is None:
+            return None
+        packed_image = Image.open(packed_image_path)
+        images = np.array_split(np.array(packed_image), 6, axis=1)
+        images = [Image.fromarray(im) for im in images]
+        return images
+
+    def maybe_upscale_image(
+        self,
+        tensor: Optional[torch.FloatTensor],
+        upscale: bool,
+        upscale_factor: int,
+        batched: bool = False,
+    ) -> Optional[torch.FloatTensor]:
+        if upscale:
+            with torch.no_grad():
+                tensor = tensor.permute(0, 3, 1, 2)
+                if batched:
+                    tensor = self.upscaler(tensor.half()).float()
+                else:
+                    tensor = torch.concat(
+                        [
+                            self.upscaler(im.unsqueeze(0).half()).float()
+                            for im in tensor
+                        ],
+                        dim=0,
+                    )
+                tensor = tensor.clamp(0, 1).permute(0, 2, 3, 1)
+            clear()
+        return tensor
+
+    def view_inpaint(
+        self,
+        mod_name: str,
+        mesh: TexturedMesh,
+        uv_proj: torch.FloatTensor,
+        uv_valid_mask: torch.BoolTensor,
+        config: ModProcessConfig,
+        debug_dir: Optional[str] = None,
+    ) -> Tuple[torch.FloatTensor, torch.BoolTensor]:
+        def inpaint_func(
+            image: torch.FloatTensor, mask: torch.FloatTensor
+        ) -> torch.FloatTensor:
+            with torch.no_grad():
+                return self.inpainter(image.to(torch.float32), mask.to(torch.float32))
+
+        return self.smart_painter(
+            mod_name,
+            mesh,
+            inpaint_func,
+            uv_proj,
+            ~uv_valid_mask,
+            max_view_score_thresh=config.view_inpaint_max_view_score_thresh,
+            min_rounds=config.view_inpaint_min_rounds,
+            max_rounds=config.view_inpaint_max_rounds,
+            uv_padding_end=config.view_inpaint_uv_padding_end,
+            debug_dir=debug_dir,
+            debug_visualize_details=False,
+        )
+
+    def __call__(
+        self,
+        mesh_path: str,
+        save_dir: str,
+        save_name: str = "default",
+        # mesh load settings
+        move_to_center: bool = False,
+        front_x: bool = True,
+        # uv unwarp
+        uv_unwarp: bool = False,
+        preprocess_mesh: bool = False,
+        # projection
+        uv_size: int = 4096,
+        # modes
+        rgb_path: Optional[str] = None,
+        rgb_process_config: ModProcessConfig = ModProcessConfig(),
+        base_color_path: Optional[str] = None,
+        base_color_process_config: ModProcessConfig = ModProcessConfig(),
+        orm_path: Optional[str] = None,
+        orm_process_config: ModProcessConfig = ModProcessConfig(),
+        normal_path: Optional[str] = None,
+        normal_strength: float = 1.0,
+        normal_process_config: ModProcessConfig = ModProcessConfig(),
+        # inpaint
+        uv_inpaint_use_network: bool = False,
+        view_inpaint_include_occlusion_boundary: bool = False,
+        poisson_reprojection: bool = False,
+        # camera
+        camera_projection_type: str = "ORTHO",
+        camera_elevation_deg: List[float] = [0, 0, 0, 0, 89.99, -89.99],
+        camera_azimuth_deg: List[float] = [0, 90, 180, 270, 180, 180],
+        camera_distance: float = 1.0,
+        camera_ortho_scale: float = 1.1,
+        camera_fov_deg: float = 40,
+        # debug
+        debug_mode: bool = False,
+        # seam removal
+        apply_dilate: bool = True,
+    ):
+        clear()
+
+        if debug_mode:
+            debug_dir = os.path.join(save_dir, "debug")
+            os.makedirs(debug_dir, exist_ok=True)
+
+        if uv_unwarp:
+            file_suffix = os.path.splitext(mesh_path)[-1]
+            mesh_path_new = mesh_path.replace(file_suffix, f"_unwarp{file_suffix}")
+            process_raw(mesh_path, mesh_path_new, preprocess=preprocess_mesh)
+            mesh_path = mesh_path_new
+
+        mesh: TexturedMesh = load_mesh(
+            mesh_path,
+            rescale=True,
+            move_to_center=move_to_center,
+            front_x_to_y=front_x,
+            default_uv_size=uv_size,
+            device=self.device,
+        )
+
+        # projection
+        if camera_projection_type == "PERSP":
+            raise NotImplementedError
+        elif camera_projection_type == "ORTHO":
+            cameras = get_orthogonal_camera(
+                elevation_deg=camera_elevation_deg,
+                distance=[camera_distance] * 6,
+                left=-camera_ortho_scale / 2,
+                right=camera_ortho_scale / 2,
+                bottom=-camera_ortho_scale / 2,
+                top=camera_ortho_scale / 2,
+                azimuth_deg=[x - 90 for x in camera_azimuth_deg],  # -y as front
+                device=self.device,
+            )
+
+        mod_kwargs = {
+            "rgb": (rgb_path, rgb_process_config),
+            "base_color": (base_color_path, base_color_process_config),
+            "orm": (orm_path, orm_process_config),
+            "normal": (normal_path, normal_process_config),
+        }
+        mod_uv_image, mod_uv_tensor = {}, {}
+        for mod_name, (mod_path, mod_process_config) in mod_kwargs.items():
+            if mod_path is None:
+                mod_uv_image[mod_name] = None
+                mod_uv_tensor[mod_name] = None
+                continue
+            mod_images = self.load_packed_images(mod_path)
+            mod_tensor = image_to_tensor(mod_images, device=self.device)
+            mod_tensor = self.maybe_upscale_image(
+                mod_tensor,
+                mod_process_config.view_upscale,
+                mod_process_config.view_upscale_factor,
+            )
+            if mod_process_config.view_upscale and debug_mode:
+                make_image_grid(tensor_to_image(mod_tensor, batched=True), rows=1).save(
+                    os.path.join(debug_dir, f"{mod_name}_upscaled.jpg")
+                )
+
+            if mod_name == "normal":
+                _, height, width, _ = mod_tensor.shape
+                render_out = render(
+                    self.ctx,
+                    mesh,
+                    cameras,
+                    height=height,
+                    width=width,
+                    render_attr=False,
+                    render_depth=False,
+                    render_normal=True,
+                    render_tangent=True,
+                )
+
+                # compute UV tangent space
+                vN = render_out.normal
+                vT = render_out.tangent
+                vB = torch.cross(vN, vT, dim=-1)
+                tang_space = F.normalize(torch.stack([vT, vB, vN], dim=-2), dim=-1)
+
+                # compute geometry tangent space
+                vGN = vN
+                vGT = torch.as_tensor(
+                    [
+                        [1, 0, 0],
+                        [0, 1, 0],
+                        [-1, 0, 0],
+                        [0, -1, 0],
+                        [-1, 0, 0],
+                        [-1, 0, 0],
+                    ],
+                    dtype=torch.float32,
+                    device=self.device,
+                )[:, None, None, :]
+                vGB = torch.cross(vGN, vGT, dim=-1)
+                vGT = torch.cross(vGB, vGN, dim=-1)
+                geo_tang_space = F.normalize(
+                    torch.stack([vGT, vGB, vGN], dim=-2), dim=-1
+                )
+
+                # restore world-space normal from geometry tangent space
+                mod_tensor = mod_tensor * 2 - 1
+                mod_tensor = F.normalize(
+                    (
+                        mod_tensor[:, :, :, None, :]
+                        * geo_tang_space.permute(0, 1, 2, 4, 3)
+                    ).sum(-1),
+                    dim=-1,
+                )
+
+                # bake world-space normal to UV tangent space
+                mod_tensor = F.normalize(
+                    (mod_tensor[:, :, :, None, :] * tang_space).sum(-1), dim=-1
+                )
+                mod_tensor = (mod_tensor * 0.5 + 0.5).clamp(0, 1)
+
+                uv_proj, uv_valid_mask = self.cam_proj(
+                    mod_tensor,
+                    mesh,
+                    cameras,
+                    from_scratch=mod_process_config.inpaint_mode != "none",
+                    poisson_blending=False,
+                    depth_grad_dilation=5,
+                    uv_exp_blend_alpha=3,
+                    uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
+                    aoi_cos_valid_threshold=0.2,
+                    uv_size=uv_size,
+                    return_uv_projection_mask=True,
+                )
+                uv_proj[~uv_valid_mask] = torch.as_tensor([0.5, 0.5, 1]).to(uv_proj)
+            else:
+                # TODO: tweak depth_grad_dilation
+                cam_proj_out = self.cam_proj(
+                    mod_tensor,
+                    mesh,
+                    cameras,
+                    from_scratch=mod_process_config.inpaint_mode != "none",
+                    poisson_blending=False,
+                    depth_grad_dilation=5,
+                    depth_grad_threshold=0.1,
+                    uv_exp_blend_alpha=3,
+                    uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
+                    aoi_cos_valid_threshold=0.2,
+                    uv_size=uv_size,
+                    uv_padding=not uv_inpaint_use_network,
+                    return_dict=True,
+                )
+                uv_proj, uv_valid_mask, uv_depth_grad = (
+                    cam_proj_out.uv_proj,
+                    cam_proj_out.uv_proj_mask,
+                    cam_proj_out.uv_depth_grad,
+                )
+                if uv_inpaint_use_network:
+                    uv_inpaint_mask_input = 1 - uv_valid_mask[None, None].float()
+                    uv_inpaint_image_input = uv_proj[None].permute(0, 3, 1, 2)
+                    with torch.no_grad():
+                        uv_inpaint_result = self.inpainter(
+                            uv_inpaint_image_input, uv_inpaint_mask_input
+                        )[0].permute(1, 2, 0)
+                    clear()
+                    if debug_mode:
+                        make_image_grid(
+                            [
+                                tensor_to_image(uv_proj),
+                                tensor_to_image(uv_valid_mask),
+                                tensor_to_image(uv_inpaint_result),
+                            ]
+                        ).save(os.path.join(debug_dir, f"{mod_name}_uv_inpaint.jpg"))
+                    uv_proj = uv_inpaint_result.contiguous()
+
+                if mod_process_config.inpaint_mode == "view":
+                    if view_inpaint_include_occlusion_boundary:
+                        uv_max_depth_grad = uv_depth_grad.max(dim=0)[0]
+                        uv_valid_mask = uv_valid_mask & (uv_max_depth_grad < 0.1)
+                    uv_proj, uv_valid_mask = self.view_inpaint(
+                        mod_name,
+                        mesh,
+                        uv_proj,
+                        uv_valid_mask,
+                        mod_process_config,
+                        debug_dir=debug_dir if debug_mode else None,
+                    )
+
+                if poisson_reprojection:
+                    # up and down
+                    mesh.texture = uv_proj
+                    uv_proj = self.cam_proj(
+                        mod_tensor[4:5],
+                        mesh,
+                        cameras[4:5],
+                        from_scratch=False,
+                        poisson_blending=True,
+                        pb_keep_original_border=True,
+                        depth_grad_dilation=5,
+                        uv_exp_blend_alpha=3,
+                        uv_exp_blend_view_weight=torch.as_tensor([1, 1]),
+                        aoi_cos_valid_threshold=0.2,
+                        uv_size=uv_size,
+                        uv_padding=True,
+                        return_dict=False,
+                    )
+                    # front, sides and back
+                    mesh.texture = uv_proj
+                    uv_proj = self.cam_proj(
+                        mod_tensor[0:4],
+                        mesh,
+                        cameras[0:4],
+                        from_scratch=False,
+                        poisson_blending=True,
+                        pb_keep_original_border=True,
+                        depth_grad_dilation=5,
+                        uv_exp_blend_alpha=3,
+                        uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1]),
+                        aoi_cos_valid_threshold=0.2,
+                        uv_size=uv_size,
+                        uv_padding=True,
+                        return_dict=False,
+                    )
+
+                if mod_name == "orm":
+                    uv_proj[:, :, 0] = 1.0
+
+            # Apply dilate to remove texture seams before final conversion
+            if mod_name in ["rgb", "base_color"] and apply_dilate:  # Apply to color textures only
+                uv_proj = torch_image_dilate(uv_proj)
+
+            mod_uv_image[mod_name] = tensor_to_image(uv_proj)
+            mod_uv_tensor[mod_name] = uv_proj
+            clear()
+
+        shaded_model_save_path = None
+        if mod_uv_image["rgb"] is not None:
+            shaded_model_save_path = os.path.join(save_dir, f"{save_name}_shaded.glb")
+            replace_mesh_texture_and_save(
+                mesh_path,
+                shaded_model_save_path,
+                texture=mod_uv_image["rgb"],
+                backend="trimesh",
+                task_id=save_name,
+            )
+        pbr_model_save_path = None
+        if mod_uv_image["base_color"] is not None:
+            pbr_model_save_path = os.path.join(save_dir, f"{save_name}_pbr.glb")
+            replace_mesh_texture_and_save(
+                mesh_path,
+                pbr_model_save_path,
+                texture=mod_uv_image["base_color"],
+                metallic_roughness_texture=mod_uv_image["orm"],
+                normal_texture=mod_uv_image["normal"],
+                normal_strength=normal_strength,
+                backend="trimesh",
+                task_id=save_name,
+            )
+
+        clear()
+
+        return TexturePipelineOutput(
+            shaded_model_save_path=shaded_model_save_path,
+            pbr_model_save_path=pbr_model_save_path,
+        )

Gen_3D_Modules/MV_Adapter/mvadapter/schedulers/scheduler_utils.py

@@ -0,0 +1,70 @@
+import torch
+
+
+def get_sigmas(noise_scheduler, timesteps, n_dim=4, dtype=torch.float32, device=None):
+    sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
+    schedule_timesteps = noise_scheduler.timesteps.to(device)
+    timesteps = timesteps.to(device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+    sigma = sigmas[step_indices].flatten()
+    while len(sigma.shape) < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma
+
+
+def SNR_to_betas(snr):
+    """
+    Converts SNR to betas
+    """
+    # alphas_cumprod = pass
+    # snr = (alpha / ) ** 2
+    # alpha_t^2 / (1 - alpha_t^2) = snr
+    alpha_t = (snr / (1 + snr)) ** 0.5
+    alphas_cumprod = alpha_t**2
+    alphas = alphas_cumprod / torch.cat(
+        [torch.ones(1, device=snr.device), alphas_cumprod[:-1]]
+    )
+    betas = 1 - alphas
+    return betas
+
+
+def compute_snr(timesteps, noise_scheduler):
+    """
+    Computes SNR as per Min-SNR-Diffusion-Training/guided_diffusion/gaussian_diffusion.py at 521b624bd70c67cee4bdf49225915f5
+    """
+    alphas_cumprod = noise_scheduler.alphas_cumprod
+    sqrt_alphas_cumprod = alphas_cumprod**0.5
+    sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
+
+    # Expand the tensors.
+    # Adapted from Min-SNR-Diffusion-Training/guided_diffusion/gaussian_diffusion.py at 521b624bd70c67cee4bdf49225915f5
+    sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
+        timesteps
+    ].float()
+    while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
+        sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
+    alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
+
+    sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(
+        device=timesteps.device
+    )[timesteps].float()
+    while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
+        sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
+    sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
+
+    # Compute SNR.
+    snr = (alpha / sigma) ** 2
+    return snr
+
+
+def compute_alpha(timesteps, noise_scheduler):
+    alphas_cumprod = noise_scheduler.alphas_cumprod
+    sqrt_alphas_cumprod = alphas_cumprod**0.5
+    sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
+        timesteps
+    ].float()
+    while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
+        sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
+    alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
+
+    return alpha

Gen_3D_Modules/MV_Adapter/mvadapter/schedulers/scheduling_shift_snr.py

@@ -0,0 +1,138 @@
+from typing import Any
+
+import torch
+
+from .scheduler_utils import SNR_to_betas, compute_snr
+
+
+class ShiftSNRScheduler:
+    def __init__(
+        self,
+        noise_scheduler: Any,
+        timesteps: Any,
+        shift_scale: float,
+        scheduler_class: Any,
+    ):
+        self.noise_scheduler = noise_scheduler
+        self.timesteps = timesteps
+        self.shift_scale = shift_scale
+        self.scheduler_class = scheduler_class
+
+    def _get_shift_scheduler(self):
+        """
+        Prepare scheduler for shifted betas.
+
+        :return: A scheduler object configured with shifted betas
+        """
+        snr = compute_snr(self.timesteps, self.noise_scheduler)
+        shifted_betas = SNR_to_betas(snr / self.shift_scale)
+
+        return self.scheduler_class.from_config(
+            self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
+        )
+
+    def _get_interpolated_shift_scheduler(self):
+        """
+        Prepare scheduler for shifted betas and interpolate with the original betas in log space.
+
+        :return: A scheduler object configured with interpolated shifted betas
+        """
+        snr = compute_snr(self.timesteps, self.noise_scheduler)
+        shifted_snr = snr / self.shift_scale
+
+        weighting = self.timesteps.float() / (
+            self.noise_scheduler.config.num_train_timesteps - 1
+        )
+        interpolated_snr = torch.exp(
+            torch.log(snr) * (1 - weighting) + torch.log(shifted_snr) * weighting
+        )
+
+        shifted_betas = SNR_to_betas(interpolated_snr)
+
+        return self.scheduler_class.from_config(
+            self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
+        )
+
+    @classmethod
+    def from_scheduler(
+        cls,
+        noise_scheduler: Any,
+        shift_mode: str = "default",
+        timesteps: Any = None,
+        shift_scale: float = 1.0,
+        scheduler_class: Any = None,
+    ):
+        # Check input
+        if timesteps is None:
+            timesteps = torch.arange(0, noise_scheduler.config.num_train_timesteps)
+        if scheduler_class is None:
+            scheduler_class = noise_scheduler.__class__
+
+        # Create scheduler
+        shift_scheduler = cls(
+            noise_scheduler=noise_scheduler,
+            timesteps=timesteps,
+            shift_scale=shift_scale,
+            scheduler_class=scheduler_class,
+        )
+
+        if shift_mode == "default":
+            return shift_scheduler._get_shift_scheduler()
+        elif shift_mode == "interpolated":
+            return shift_scheduler._get_interpolated_shift_scheduler()
+        else:
+            raise ValueError(f"Unknown shift_mode: {shift_mode}")
+
+
+if __name__ == "__main__":
+    """
+    Compare the alpha values for different noise schedulers.
+    """
+    import matplotlib.pyplot as plt
+    from diffusers import DDPMScheduler
+
+    from .scheduler_utils import compute_alpha
+
+    # Base
+    timesteps = torch.arange(0, 1000)
+    noise_scheduler_base = DDPMScheduler.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", subfolder="scheduler"
+    )
+    alpha = compute_alpha(timesteps, noise_scheduler_base)
+    plt.plot(timesteps.numpy(), alpha.numpy(), label="Base")
+
+    # Kolors
+    num_train_timesteps_ = 1100
+    timesteps_ = torch.arange(0, num_train_timesteps_)
+    noise_kwargs = {"beta_end": 0.014, "num_train_timesteps": num_train_timesteps_}
+    noise_scheduler_kolors = DDPMScheduler.from_config(
+        noise_scheduler_base.config, **noise_kwargs
+    )
+    alpha = compute_alpha(timesteps_, noise_scheduler_kolors)
+    plt.plot(timesteps_.numpy(), alpha.numpy(), label="Kolors")
+
+    # Shift betas
+    shift_scale = 8.0
+    noise_scheduler_shift = ShiftSNRScheduler.from_scheduler(
+        noise_scheduler_base, shift_mode="default", shift_scale=shift_scale
+    )
+    alpha = compute_alpha(timesteps, noise_scheduler_shift)
+    plt.plot(timesteps.numpy(), alpha.numpy(), label="Shift Noise (scale 8.0)")
+
+    # Shift betas (interpolated)
+    noise_scheduler_inter = ShiftSNRScheduler.from_scheduler(
+        noise_scheduler_base, shift_mode="interpolated", shift_scale=shift_scale
+    )
+    alpha = compute_alpha(timesteps, noise_scheduler_inter)
+    plt.plot(timesteps.numpy(), alpha.numpy(), label="Interpolated (scale 8.0)")
+
+    # ZeroSNR
+    noise_scheduler = DDPMScheduler.from_config(
+        noise_scheduler_base.config, rescale_betas_zero_snr=True
+    )
+    alpha = compute_alpha(timesteps, noise_scheduler)
+    plt.plot(timesteps.numpy(), alpha.numpy(), label="ZeroSNR")
+
+    plt.legend()
+    plt.grid()
+    plt.savefig("check_alpha.png")

Gen_3D_Modules/MV_Adapter/mvadapter/utils/__init__.py

@@ -0,0 +1,3 @@
+from .mesh_utils.camera import get_camera, get_orthogonal_camera
+from .geometry import get_plucker_embeds_from_cameras_ortho
+from .saving import make_image_grid, tensor_to_image, image_to_tensor

Gen_3D_Modules/MV_Adapter/mvadapter/utils/base.py

@@ -0,0 +1,123 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from .config import parse_structured
+from .misc import get_device, load_module_weights
+from .typing import *
+
+
+class Configurable:
+    @dataclass
+    class Config:
+        pass
+
+    def __init__(self, cfg: Optional[dict] = None) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+
+
+class Updateable:
+    def do_update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ):
+        for attr in self.__dir__():
+            if attr.startswith("_"):
+                continue
+            try:
+                module = getattr(self, attr)
+            except:
+                continue  # ignore attributes like property, which can't be retrived using getattr?
+            if isinstance(module, Updateable):
+                module.do_update_step(
+                    epoch, global_step, on_load_weights=on_load_weights
+                )
+        self.update_step(epoch, global_step, on_load_weights=on_load_weights)
+
+    def do_update_step_end(self, epoch: int, global_step: int):
+        for attr in self.__dir__():
+            if attr.startswith("_"):
+                continue
+            try:
+                module = getattr(self, attr)
+            except:
+                continue  # ignore attributes like property, which can't be retrived using getattr?
+            if isinstance(module, Updateable):
+                module.do_update_step_end(epoch, global_step)
+        self.update_step_end(epoch, global_step)
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # override this method to implement custom update logic
+        # if on_load_weights is True, you should be careful doing things related to model evaluations,
+        # as the models and tensors are not guarenteed to be on the same device
+        pass
+
+    def update_step_end(self, epoch: int, global_step: int):
+        pass
+
+
+def update_if_possible(module: Any, epoch: int, global_step: int) -> None:
+    if isinstance(module, Updateable):
+        module.do_update_step(epoch, global_step)
+
+
+def update_end_if_possible(module: Any, epoch: int, global_step: int) -> None:
+    if isinstance(module, Updateable):
+        module.do_update_step_end(epoch, global_step)
+
+
+class BaseObject(Updateable):
+    @dataclass
+    class Config:
+        pass
+
+    cfg: Config  # add this to every subclass of BaseObject to enable static type checking
+
+    def __init__(
+        self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+        self.device = get_device()
+        self.configure(*args, **kwargs)
+
+    def configure(self, *args, **kwargs) -> None:
+        pass
+
+
+class BaseModule(nn.Module, Updateable):
+    @dataclass
+    class Config:
+        weights: Optional[str] = None
+
+    cfg: Config  # add this to every subclass of BaseModule to enable static type checking
+
+    def __init__(
+        self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+        self.device = get_device()
+        self._non_modules = {}
+        self.configure(*args, **kwargs)
+        if self.cfg.weights is not None:
+            # format: path/to/weights:module_name
+            weights_path, module_name = self.cfg.weights.split(":")
+            state_dict, epoch, global_step = load_module_weights(
+                weights_path, module_name=module_name, map_location="cpu"
+            )
+            self.load_state_dict(state_dict)
+            self.do_update_step(
+                epoch, global_step, on_load_weights=True
+            )  # restore states
+
+    def configure(self, *args, **kwargs) -> None:
+        pass
+
+    def register_non_module(self, name: str, module: nn.Module) -> None:
+        # non-modules won't be treated as model parameters
+        self._non_modules[name] = module
+
+    def non_module(self, name: str):
+        return self._non_modules.get(name, None)

Gen_3D_Modules/MV_Adapter/mvadapter/utils/callbacks.py

@@ -0,0 +1,158 @@
+import os
+import shutil
+import subprocess
+
+import pytorch_lightning
+
+from .config import dump_config
+from .misc import parse_version
+
+if parse_version(pytorch_lightning.__version__) > parse_version("1.8"):
+    from pytorch_lightning.callbacks import Callback
+else:
+    from pytorch_lightning.callbacks.base import Callback
+
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
+from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
+
+
+class VersionedCallback(Callback):
+    def __init__(self, save_root, version=None, use_version=True):
+        self.save_root = save_root
+        self._version = version
+        self.use_version = use_version
+
+    @property
+    def version(self) -> int:
+        """Get the experiment version.
+
+        Returns:
+            The experiment version if specified else the next version.
+        """
+        if self._version is None:
+            self._version = self._get_next_version()
+        return self._version
+
+    def _get_next_version(self):
+        existing_versions = []
+        if os.path.isdir(self.save_root):
+            for f in os.listdir(self.save_root):
+                bn = os.path.basename(f)
+                if bn.startswith("version_"):
+                    dir_ver = os.path.splitext(bn)[0].split("_")[1].replace("/", "")
+                    existing_versions.append(int(dir_ver))
+        if len(existing_versions) == 0:
+            return 0
+        return max(existing_versions) + 1
+
+    @property
+    def savedir(self):
+        if not self.use_version:
+            return self.save_root
+        return os.path.join(
+            self.save_root,
+            (
+                self.version
+                if isinstance(self.version, str)
+                else f"version_{self.version}"
+            ),
+        )
+
+
+class CodeSnapshotCallback(VersionedCallback):
+    def __init__(self, save_root, version=None, use_version=True):
+        super().__init__(save_root, version, use_version)
+
+    def get_file_list(self):
+        return [
+            b.decode()
+            for b in set(
+                subprocess.check_output(
+                    'git ls-files -- ":!:load/*"', shell=True
+                ).splitlines()
+            )
+            | set(  # hard code, TODO: use config to exclude folders or files
+                subprocess.check_output(
+                    "git ls-files --others --exclude-standard", shell=True
+                ).splitlines()
+            )
+        ]
+
+    @rank_zero_only
+    def save_code_snapshot(self):
+        os.makedirs(self.savedir, exist_ok=True)
+        for f in self.get_file_list():
+            if not os.path.exists(f) or os.path.isdir(f):
+                continue
+            os.makedirs(os.path.join(self.savedir, os.path.dirname(f)), exist_ok=True)
+            shutil.copyfile(f, os.path.join(self.savedir, f))
+
+    def on_fit_start(self, trainer, pl_module):
+        try:
+            self.save_code_snapshot()
+        except:
+            rank_zero_warn(
+                "Code snapshot is not saved. Please make sure you have git installed and are in a git repository."
+            )
+
+
+class ConfigSnapshotCallback(VersionedCallback):
+    def __init__(self, config_path, config, save_root, version=None, use_version=True):
+        super().__init__(save_root, version, use_version)
+        self.config_path = config_path
+        self.config = config
+
+    @rank_zero_only
+    def save_config_snapshot(self):
+        os.makedirs(self.savedir, exist_ok=True)
+        dump_config(os.path.join(self.savedir, "parsed.yaml"), self.config)
+        shutil.copyfile(self.config_path, os.path.join(self.savedir, "raw.yaml"))
+
+    def on_fit_start(self, trainer, pl_module):
+        self.save_config_snapshot()
+
+
+class CustomProgressBar(TQDMProgressBar):
+    def get_metrics(self, *args, **kwargs):
+        # don't show the version number
+        items = super().get_metrics(*args, **kwargs)
+        items.pop("v_num", None)
+        return items
+
+
+class ProgressCallback(Callback):
+    def __init__(self, save_path):
+        super().__init__()
+        self.save_path = save_path
+        self._file_handle = None
+
+    @property
+    def file_handle(self):
+        if self._file_handle is None:
+            self._file_handle = open(self.save_path, "w")
+        return self._file_handle
+
+    @rank_zero_only
+    def write(self, msg: str) -> None:
+        self.file_handle.seek(0)
+        self.file_handle.truncate()
+        self.file_handle.write(msg)
+        self.file_handle.flush()
+
+    @rank_zero_only
+    def on_train_batch_end(self, trainer, pl_module, *args, **kwargs):
+        self.write(
+            f"Generation progress: {pl_module.true_global_step / trainer.max_steps * 100:.2f}%"
+        )
+
+    @rank_zero_only
+    def on_validation_start(self, trainer, pl_module):
+        self.write(f"Rendering validation image ...")
+
+    @rank_zero_only
+    def on_test_start(self, trainer, pl_module):
+        self.write(f"Rendering video ...")
+
+    @rank_zero_only
+    def on_predict_start(self, trainer, pl_module):
+        self.write(f"Exporting mesh assets ...")

Gen_3D_Modules/MV_Adapter/mvadapter/utils/config.py

@@ -0,0 +1,140 @@
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+
+from omegaconf import OmegaConf
+
+from .core import debug, find, info, warn
+from .typing import *
+
+# ============ Register OmegaConf Resolvers ============= #
+OmegaConf.register_new_resolver(
+    "calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n)
+)
+OmegaConf.register_new_resolver("add", lambda a, b: a + b)
+OmegaConf.register_new_resolver("sub", lambda a, b: a - b)
+OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
+OmegaConf.register_new_resolver("div", lambda a, b: a / b)
+OmegaConf.register_new_resolver("idiv", lambda a, b: a // b)
+OmegaConf.register_new_resolver("basename", lambda p: os.path.basename(p))
+OmegaConf.register_new_resolver("rmspace", lambda s, sub: s.replace(" ", sub))
+OmegaConf.register_new_resolver("tuple2", lambda s: [float(s), float(s)])
+OmegaConf.register_new_resolver("gt0", lambda s: s > 0)
+OmegaConf.register_new_resolver("not", lambda s: not s)
+
+
+def calc_num_train_steps(num_data, batch_size, max_epochs, num_nodes, num_cards=8):
+    return int(num_data / (num_nodes * num_cards * batch_size)) * max_epochs
+
+
+OmegaConf.register_new_resolver("calc_num_train_steps", calc_num_train_steps)
+
+# ======================================================= #
+
+
+# ============== Automatic Name Resolvers =============== #
+def get_naming_convention(cfg):
+    # TODO
+    name = f"lrm_{cfg.system.backbone.num_layers}"
+    return name
+
+
+# ======================================================= #
+
+
+@dataclass
+class ExperimentConfig:
+    name: str = "default"
+    description: str = ""
+    tag: str = ""
+    seed: int = 0
+    use_timestamp: bool = True
+    timestamp: Optional[str] = None
+    exp_root_dir: str = "outputs"
+
+    ### these shouldn't be set manually
+    exp_dir: str = "outputs/default"
+    trial_name: str = "exp"
+    trial_dir: str = "outputs/default/exp"
+    n_gpus: int = 1
+    ###
+
+    resume: Optional[str] = None
+
+    data_cls: str = ""
+    data: dict = field(default_factory=dict)
+
+    system_cls: str = ""
+    system: dict = field(default_factory=dict)
+
+    # accept pytorch-lightning trainer parameters
+    # see https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api
+    trainer: dict = field(default_factory=dict)
+
+    # accept pytorch-lightning checkpoint callback parameters
+    # see https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
+    checkpoint: dict = field(default_factory=dict)
+
+
+def load_config(
+    *yamls: str, cli_args: list = [], from_string=False, makedirs=True, **kwargs
+) -> Any:
+    if from_string:
+        parse_func = OmegaConf.create
+    else:
+        parse_func = OmegaConf.load
+    yaml_confs = []
+    for y in yamls:
+        conf = parse_func(y)
+        extends = conf.pop("extends", None)
+        if extends:
+            assert os.path.exists(extends), f"File {extends} does not exist."
+            yaml_confs.append(OmegaConf.load(extends))
+        yaml_confs.append(conf)
+    cli_conf = OmegaConf.from_cli(cli_args)
+    cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs)
+    OmegaConf.resolve(cfg)
+    assert isinstance(cfg, DictConfig)
+    scfg: ExperimentConfig = parse_structured(ExperimentConfig, cfg)
+
+    # post processing
+    # auto naming
+    if scfg.name == "auto":
+        scfg.name = get_naming_convention(scfg)
+    # add timestamp
+    if not scfg.tag and not scfg.use_timestamp:
+        raise ValueError("Either tag is specified or use_timestamp is True.")
+    scfg.trial_name = scfg.tag
+    # if resume from an existing config, scfg.timestamp should not be None
+    if scfg.timestamp is None:
+        scfg.timestamp = ""
+        if scfg.use_timestamp:
+            if scfg.n_gpus > 1:
+                warn(
+                    "Timestamp is disabled when using multiple GPUs, please make sure you have a unique tag."
+                )
+            else:
+                scfg.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S")
+    # make directories
+    scfg.trial_name += scfg.timestamp
+    scfg.exp_dir = os.path.join(scfg.exp_root_dir, scfg.name)
+    scfg.trial_dir = os.path.join(scfg.exp_dir, scfg.trial_name)
+
+    if makedirs:
+        os.makedirs(scfg.trial_dir, exist_ok=True)
+
+    return scfg
+
+
+def config_to_primitive(config, resolve: bool = True) -> Any:
+    return OmegaConf.to_container(config, resolve=resolve)
+
+
+def dump_config(path: str, config) -> None:
+    with open(path, "w") as fp:
+        OmegaConf.save(config=config, f=fp)
+
+
+def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
+    scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg)
+    return scfg

Gen_3D_Modules/MV_Adapter/mvadapter/utils/core.py

@@ -0,0 +1,29 @@
+import importlib
+
+###  grammar sugar for logging utilities  ###
+import logging
+
+logger = logging.getLogger("pytorch_lightning")
+
+from pytorch_lightning.utilities.rank_zero import (
+    rank_zero_debug,
+    rank_zero_info,
+    rank_zero_only,
+)
+
+
+def find(cls_string):
+    module_string = ".".join(cls_string.split(".")[:-1])
+    cls_name = cls_string.split(".")[-1]
+    module = importlib.import_module(module_string, package=None)
+    cls = getattr(module, cls_name)
+    return cls
+
+
+debug = rank_zero_debug
+info = rank_zero_info
+
+
+@rank_zero_only
+def warn(*args, **kwargs):
+    logger.warn(*args, **kwargs)

Gen_3D_Modules/MV_Adapter/mvadapter/utils/geometry.py

@@ -0,0 +1,253 @@
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+def get_position_map_from_depth(depth, mask, intrinsics, extrinsics, image_wh=None):
+    """Compute the position map from the depth map and the camera parameters for a batch of views.
+
+    Args:
+        depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
+        mask (torch.Tensor): The masks with the shape (B, H, W, 1).
+        intrinsics (torch.Tensor): The camera intrinsics matrices with the shape (B, 3, 3).
+        extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
+        image_wh (Tuple[int, int]): The image width and height.
+
+    Returns:
+        torch.Tensor: The position maps with the shape (B, H, W, 3).
+    """
+    if image_wh is None:
+        image_wh = depth.shape[2], depth.shape[1]
+
+    B, H, W, _ = depth.shape
+    depth = depth.squeeze(-1)
+
+    u_coord, v_coord = torch.meshgrid(
+        torch.arange(image_wh[0]), torch.arange(image_wh[1]), indexing="xy"
+    )
+    u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
+    v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
+
+    # Compute the position map by back-projecting depth pixels to 3D space
+    x = (
+        (u_coord - intrinsics[:, 0, 2].unsqueeze(-1).unsqueeze(-1))
+        * depth
+        / intrinsics[:, 0, 0].unsqueeze(-1).unsqueeze(-1)
+    )
+    y = (
+        (v_coord - intrinsics[:, 1, 2].unsqueeze(-1).unsqueeze(-1))
+        * depth
+        / intrinsics[:, 1, 1].unsqueeze(-1).unsqueeze(-1)
+    )
+    z = depth
+
+    # Concatenate to form the 3D coordinates in the camera frame
+    camera_coords = torch.stack([x, y, z], dim=-1)
+
+    # Apply the extrinsic matrix to get coordinates in the world frame
+    coords_homogeneous = torch.nn.functional.pad(
+        camera_coords, (0, 1), "constant", 1.0
+    )  # Add a homogeneous coordinate
+    world_coords = torch.matmul(
+        coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
+    ).view(B, H, W, 4)
+
+    # Apply the mask to the position map
+    position_map = world_coords[..., :3] * mask
+
+    return position_map
+
+
+def get_position_map_from_depth_ortho(
+    depth, mask, extrinsics, ortho_scale, image_wh=None
+):
+    """Compute the position map from the depth map and the camera parameters for a batch of views
+    using orthographic projection with a given ortho_scale.
+
+    Args:
+        depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
+        mask (torch.Tensor): The masks with the shape (B, H, W, 1).
+        extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
+        ortho_scale (torch.Tensor): The scaling factor for the orthographic projection with the shape (B, 1, 1, 1).
+        image_wh (Tuple[int, int]): Optional. The image width and height.
+
+    Returns:
+        torch.Tensor: The position maps with the shape (B, H, W, 3).
+    """
+    if image_wh is None:
+        image_wh = depth.shape[2], depth.shape[1]
+
+    B, H, W, _ = depth.shape
+    depth = depth.squeeze(-1)
+
+    # Generating grid of coordinates in the image space
+    u_coord, v_coord = torch.meshgrid(
+        torch.arange(0, image_wh[0]), torch.arange(0, image_wh[1]), indexing="xy"
+    )
+    u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
+    v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
+
+    # Compute the position map using orthographic projection with ortho_scale
+    x = (u_coord - image_wh[0] / 2) * ortho_scale / image_wh[0]
+    y = (v_coord - image_wh[1] / 2) * ortho_scale / image_wh[1]
+    z = depth
+
+    # Concatenate to form the 3D coordinates in the camera frame
+    camera_coords = torch.stack([x, y, z], dim=-1)
+
+    # Apply the extrinsic matrix to get coordinates in the world frame
+    coords_homogeneous = torch.nn.functional.pad(
+        camera_coords, (0, 1), "constant", 1.0
+    )  # Add a homogeneous coordinate
+    world_coords = torch.matmul(
+        coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
+    ).view(B, H, W, 4)
+
+    # Apply the mask to the position map
+    position_map = world_coords[..., :3] * mask
+
+    return position_map
+
+
+def get_opencv_from_blender(matrix_world, fov=None, image_size=None):
+    # convert matrix_world to opencv format extrinsics
+    opencv_world_to_cam = matrix_world.inverse()
+    opencv_world_to_cam[1, :] *= -1
+    opencv_world_to_cam[2, :] *= -1
+    R, T = opencv_world_to_cam[:3, :3], opencv_world_to_cam[:3, 3]
+
+    if fov is None:  # orthographic camera
+        return R, T
+
+    R, T = R.unsqueeze(0), T.unsqueeze(0)
+    # convert fov to opencv format intrinsics
+    focal = 1 / np.tan(fov / 2)
+    intrinsics = np.diag(np.array([focal, focal, 1])).astype(np.float32)
+    opencv_cam_matrix = (
+        torch.from_numpy(intrinsics).unsqueeze(0).float().to(matrix_world.device)
+    )
+    opencv_cam_matrix[:, :2, -1] += torch.tensor([image_size / 2, image_size / 2]).to(
+        matrix_world.device
+    )
+    opencv_cam_matrix[:, [0, 1], [0, 1]] *= image_size / 2
+
+    return R, T, opencv_cam_matrix
+
+
+def get_ray_directions(
+    H: int,
+    W: int,
+    focal: float,
+    principal: Optional[Tuple[float, float]] = None,
+    use_pixel_centers: bool = True,
+) -> torch.Tensor:
+    """
+    Get ray directions for all pixels in camera coordinate.
+    Args:
+        H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
+    Outputs:
+        directions: (H, W, 3), the direction of the rays in camera coordinate
+    """
+    pixel_center = 0.5 if use_pixel_centers else 0
+    cx, cy = W / 2, H / 2 if principal is None else principal
+    i, j = torch.meshgrid(
+        torch.arange(W, dtype=torch.float32) + pixel_center,
+        torch.arange(H, dtype=torch.float32) + pixel_center,
+        indexing="xy",
+    )
+    directions = torch.stack(
+        [(i - cx) / focal, -(j - cy) / focal, -torch.ones_like(i)], -1
+    )
+    return F.normalize(directions, dim=-1)
+
+
+def get_rays(
+    directions: torch.Tensor, c2w: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Get ray origins and directions from camera coordinates to world coordinates
+    Args:
+        directions: (H, W, 3) ray directions in camera coordinates
+        c2w: (4, 4) camera-to-world transformation matrix
+    Outputs:
+        rays_o, rays_d: (H, W, 3) ray origins and directions in world coordinates
+    """
+    # Rotate ray directions from camera coordinate to the world coordinate
+    rays_d = directions @ c2w[:3, :3].T
+    rays_o = c2w[:3, 3].expand(rays_d.shape)
+    return rays_o, rays_d
+
+
+def compute_plucker_embed(
+    c2w: torch.Tensor, image_width: int, image_height: int, focal: float
+) -> torch.Tensor:
+    """
+    Computes Plucker coordinates for a camera.
+    Args:
+        c2w: (4, 4) camera-to-world transformation matrix
+        image_width: Image width
+        image_height: Image height
+        focal: Focal length of the camera
+    Returns:
+        plucker: (6, H, W) Plucker embedding
+    """
+    directions = get_ray_directions(image_height, image_width, focal)
+    rays_o, rays_d = get_rays(directions, c2w)
+    # Cross product to get Plucker coordinates
+    cross = torch.cross(rays_o, rays_d, dim=-1)
+    plucker = torch.cat((rays_d, cross), dim=-1)
+    return plucker.permute(2, 0, 1)
+
+
+def get_plucker_embeds_from_cameras(
+    c2w: List[torch.Tensor], fov: List[float], image_size: int
+) -> torch.Tensor:
+    """
+    Given lists of camera transformations and fov, returns the batched plucker embeddings.
+    Args:
+        c2w: list of camera-to-world transformation matrices
+        fov: list of field of view values
+        image_size: size of the image
+    Returns:
+        plucker_embeds: (B, 6, H, W) batched plucker embeddings
+    """
+    plucker_embeds = []
+    for cam_matrix, cam_fov in zip(c2w, fov):
+        focal = 0.5 * image_size / np.tan(0.5 * cam_fov)
+        plucker = compute_plucker_embed(cam_matrix, image_size, image_size, focal)
+        plucker_embeds.append(plucker)
+    return torch.stack(plucker_embeds)
+
+
+def get_plucker_embeds_from_cameras_ortho(
+    c2w: List[torch.Tensor], ortho_scale: List[float], image_size: int
+):
+    """
+    Given lists of camera transformations and fov, returns the batched plucker embeddings.
+
+    Parameters:
+        c2w: list of camera-to-world transformation matrices
+        fov: list of field of view values
+        image_size: size of the image
+
+    Returns:
+        plucker_embeds: plucker embeddings (B, 6, H, W)
+    """
+    plucker_embeds = []
+    # compute pairwise mask and plucker embeddings
+    for cam_matrix, scale in zip(c2w, ortho_scale):
+        # blender to opencv to pytorch3d
+        R, T = get_opencv_from_blender(cam_matrix)
+        cam_pos = -R.T @ T
+        view_dir = R.T @ torch.tensor([0, 0, 1]).float().to(cam_matrix.device)
+        # normalize camera position
+        cam_pos = F.normalize(cam_pos, dim=0)
+        plucker = torch.concat([view_dir, cam_pos])
+        plucker = plucker.unsqueeze(-1).unsqueeze(-1).repeat(1, image_size, image_size)
+        plucker_embeds.append(plucker)
+
+    plucker_embeds = torch.stack(plucker_embeds)
+
+    return plucker_embeds

Gen_3D_Modules/MV_Adapter/mvadapter/utils/logging.py

@@ -0,0 +1,340 @@
+# coding=utf-8
+# Copyright 2024 Optuna, Hugging Face
+#
+# 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
+#
+#     http://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.
+""" Logging utilities."""
+
+import logging
+import os
+import sys
+import threading
+from logging import CRITICAL  # NOQA
+from logging import DEBUG  # NOQA
+from logging import ERROR  # NOQA
+from logging import FATAL  # NOQA
+from logging import INFO  # NOQA
+from logging import NOTSET  # NOQA
+from logging import WARN  # NOQA
+from logging import WARNING  # NOQA
+from typing import Dict, Optional
+
+from tqdm import auto as tqdm_lib
+
+_lock = threading.Lock()
+_default_handler: Optional[logging.Handler] = None
+
+log_levels = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.INFO
+
+_tqdm_active = True
+
+
+def _get_default_logging_level() -> int:
+    """
+    If LATEXTURE_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+    not - fall back to `_default_log_level`
+    """
+    env_level_str = os.getenv("LATEXTURE_VERBOSITY", None)
+    if env_level_str:
+        if env_level_str in log_levels:
+            return log_levels[env_level_str]
+        else:
+            logging.getLogger().warning(
+                f"Unknown option LATEXTURE_VERBOSITY={env_level_str}, "
+                f"has to be one of: { ', '.join(log_levels.keys()) }"
+            )
+    return _default_log_level
+
+
+def _get_library_name() -> str:
+    return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+    return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if _default_handler:
+            # This library has already configured the library root logger.
+            return
+        _default_handler = logging.StreamHandler()  # Set sys.stderr as stream.
+        _default_handler.flush = sys.stderr.flush
+
+        # Apply our default configuration to the library root logger.
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.addHandler(_default_handler)
+        library_root_logger.setLevel(_get_default_logging_level())
+        library_root_logger.propagate = False
+
+        enable_explicit_format()
+
+
+def _reset_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if not _default_handler:
+            return
+
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.removeHandler(_default_handler)
+        library_root_logger.setLevel(logging.NOTSET)
+        _default_handler = None
+
+
+def get_log_levels_dict() -> Dict[str, int]:
+    return log_levels
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Return a logger with the specified name.
+
+    This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
+    """
+
+    if name is None:
+        name = _get_library_name()
+
+    _configure_library_root_logger()
+    return logging.getLogger(name)
+
+
+def get_verbosity() -> int:
+    """
+    Return the current level for the 🤗 Diffusers' root logger as an `int`.
+
+    Returns:
+        `int`:
+            Logging level integers which can be one of:
+
+            - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+            - `40`: `diffusers.logging.ERROR`
+            - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+            - `20`: `diffusers.logging.INFO`
+            - `10`: `diffusers.logging.DEBUG`
+
+    """
+
+    _configure_library_root_logger()
+    return _get_library_root_logger().getEffectiveLevel()
+
+
+def set_verbosity(verbosity: int) -> None:
+    """
+    Set the verbosity level for the 🤗 Diffusers' root logger.
+
+    Args:
+        verbosity (`int`):
+            Logging level which can be one of:
+
+            - `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+            - `diffusers.logging.ERROR`
+            - `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+            - `diffusers.logging.INFO`
+            - `diffusers.logging.DEBUG`
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().setLevel(verbosity)
+
+
+def set_verbosity_info() -> None:
+    """Set the verbosity to the `INFO` level."""
+    return set_verbosity(INFO)
+
+
+def set_verbosity_warning() -> None:
+    """Set the verbosity to the `WARNING` level."""
+    return set_verbosity(WARNING)
+
+
+def set_verbosity_debug() -> None:
+    """Set the verbosity to the `DEBUG` level."""
+    return set_verbosity(DEBUG)
+
+
+def set_verbosity_error() -> None:
+    """Set the verbosity to the `ERROR` level."""
+    return set_verbosity(ERROR)
+
+
+def disable_default_handler() -> None:
+    """Disable the default handler of the 🤗 Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().removeHandler(_default_handler)
+
+
+def enable_default_handler() -> None:
+    """Enable the default handler of the 🤗 Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().addHandler(_default_handler)
+
+
+def add_handler(handler: logging.Handler) -> None:
+    """adds a handler to the HuggingFace Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None
+    _get_library_root_logger().addHandler(handler)
+
+
+def remove_handler(handler: logging.Handler) -> None:
+    """removes given handler from the HuggingFace Diffusers' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None and handler in _get_library_root_logger().handlers
+    _get_library_root_logger().removeHandler(handler)
+
+
+def disable_propagation() -> None:
+    """
+    Disable propagation of the library log outputs. Note that log propagation is disabled by default.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = False
+
+
+def enable_propagation() -> None:
+    """
+    Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent
+    double logging if the root logger has been configured.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = True
+
+
+def enable_explicit_format() -> None:
+    """
+    Enable explicit formatting for every 🤗 Diffusers' logger. The explicit formatter is as follows:
+    ```
+    [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
+    ```
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        formatter = logging.Formatter(
+            "[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s"
+        )
+        handler.setFormatter(formatter)
+
+
+def reset_format() -> None:
+    """
+    Resets the formatting for 🤗 Diffusers' loggers.
+
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        handler.setFormatter(None)
+
+
+def warning_advice(self, *args, **kwargs) -> None:
+    """
+    This method is identical to `logger.warning()`, but if env var LATEXTURE_NO_ADVISORY_WARNINGS=1 is set, this
+    warning will not be printed
+    """
+    no_advisory_warnings = os.getenv("LATEXTURE_NO_ADVISORY_WARNINGS", False)
+    if no_advisory_warnings:
+        return
+    self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_advice = warning_advice
+
+
+class EmptyTqdm:
+    """Dummy tqdm which doesn't do anything."""
+
+    def __init__(self, *args, **kwargs):  # pylint: disable=unused-argument
+        self._iterator = args[0] if args else None
+
+    def __iter__(self):
+        return iter(self._iterator)
+
+    def __getattr__(self, _):
+        """Return empty function."""
+
+        def empty_fn(*args, **kwargs):  # pylint: disable=unused-argument
+            return
+
+        return empty_fn
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        return
+
+
+class _tqdm_cls:
+    def __call__(self, *args, **kwargs):
+        if _tqdm_active:
+            return tqdm_lib.tqdm(*args, **kwargs)
+        else:
+            return EmptyTqdm(*args, **kwargs)
+
+    def set_lock(self, *args, **kwargs):
+        self._lock = None
+        if _tqdm_active:
+            return tqdm_lib.tqdm.set_lock(*args, **kwargs)
+
+    def get_lock(self):
+        if _tqdm_active:
+            return tqdm_lib.tqdm.get_lock()
+
+
+tqdm = _tqdm_cls()
+
+
+def is_progress_bar_enabled() -> bool:
+    """Return a boolean indicating whether tqdm progress bars are enabled."""
+    global _tqdm_active
+    return bool(_tqdm_active)
+
+
+def enable_progress_bar() -> None:
+    """Enable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = True
+
+
+def disable_progress_bar() -> None:
+    """Disable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = False

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/__init__.py

@@ -0,0 +1,20 @@
+from .camera import (
+    Camera,
+    get_c2w,
+    get_camera,
+    get_orthogonal_camera,
+    get_orthogonal_projection_matrix,
+    get_projection_matrix,
+)
+from .mesh import TexturedMesh, load_mesh, replace_mesh_texture_and_save
+from .projection import CameraProjection, CameraProjectionOutput
+from .render import (
+    DepthControlNetNormalization,
+    DepthNormalizationStrategy,
+    NVDiffRastContextWrapper,
+    RenderOutput,
+    SimpleNormalization,
+    Zero123PlusPlusNormalization,
+    render,
+)
+from .smart_paint import SmartPainter

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/blend.py

@@ -0,0 +1,324 @@
+import time
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn.functional as F
+from torch.utils.cpp_extension import load_inline
+
+from .utils import SINGLE_IMAGE_TYPE, image_to_tensor
+
+pb_solver_cpp_source = """
+#include <torch/extension.h>
+
+#include <vector>
+
+// CUDA forward declarations
+
+void pb_solver_run_cuda(
+    torch::Tensor A,
+    torch::Tensor X,
+    torch::Tensor B,
+    torch::Tensor Xbuf,
+    int num_iters
+);
+
+// C++ interface
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+void pb_solver_run(
+    torch::Tensor A,
+    torch::Tensor X,
+    torch::Tensor B,
+    torch::Tensor Xbuf,
+    int num_iters
+) {
+    CHECK_INPUT(A);
+    CHECK_INPUT(X);
+    CHECK_INPUT(B);
+    CHECK_INPUT(Xbuf);
+
+    pb_solver_run_cuda(A, X, B, Xbuf, num_iters);
+    return;
+}
+
+"""
+
+pb_solver_cuda_source = """
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+
+
+__global__ void pb_solver_run_cuda_kernel(
+    torch::PackedTensorAccessor32<int64_t,2,torch::RestrictPtrTraits> A,
+    torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> X,
+    torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> B,
+    torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> Xbuf
+) {
+    const int index = blockIdx.x * blockDim.x + threadIdx.x;
+    
+    if (index < A.size(0)) {
+        Xbuf[index] = (X[A[index][0]] + X[A[index][1]] + X[A[index][2]] + X[A[index][3]] + B[index]) / 4.;
+    }
+}
+
+void pb_solver_run_cuda(
+    torch::Tensor A,
+    torch::Tensor X,
+    torch::Tensor B,
+    torch::Tensor Xbuf,
+    int num_iters
+) {
+    int batch_size = A.size(0);
+
+    const int threads = 1024;
+    const dim3 blocks((batch_size + threads - 1) / threads);
+    
+    auto A_ptr = A.packed_accessor32<int64_t,2,torch::RestrictPtrTraits>();
+    auto X_ptr = X.packed_accessor32<float,1,torch::RestrictPtrTraits>();
+    auto B_ptr = B.packed_accessor32<float,1,torch::RestrictPtrTraits>();
+    auto Xbuf_ptr = Xbuf.packed_accessor32<float,1,torch::RestrictPtrTraits>();
+
+    for (int i = 0; i < num_iters; ++i) {
+        pb_solver_run_cuda_kernel<<<blocks, threads>>>(
+            A_ptr,
+            X_ptr,
+            B_ptr,
+            Xbuf_ptr
+        );
+        cudaDeviceSynchronize();
+        std::swap(X_ptr, Xbuf_ptr);
+    }
+    // we may waste an iteration here, but it's fine
+    return;
+}
+"""
+
+
+class PBBackend(ABC):
+    def __init__(self, *args, **kwargs) -> None:
+        pass
+
+    def solve(self, num_iters, A, X, B, Xbuf) -> None:
+        pass
+
+
+try:
+    import triton
+    import triton.language as tl
+
+    @triton.jit
+    def pb_triton_step_kernel(
+        A_ptr,
+        X_ptr,
+        B_ptr,
+        Xbuf_ptr,
+        A_row_stride: tl.constexpr,
+        n_elements: tl.constexpr,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+        block_start = pid * BLOCK_SIZE
+        offsets = tl.arange(0, BLOCK_SIZE) + block_start
+
+        A_start_ptr = A_ptr + block_start * A_row_stride
+        A_ptrs = A_start_ptr + (
+            tl.arange(0, BLOCK_SIZE)[:, None] * A_row_stride
+            + tl.arange(0, A_row_stride)[None, :]
+        )
+        B_ptrs = B_ptr + offsets
+        Xbuf_ptrs = Xbuf_ptr + offsets
+
+        mask = offsets < n_elements
+
+        A = tl.load(A_ptrs, mask=mask[:, None], other=0)
+        X = tl.load(X_ptr + A)
+        B = tl.load(B_ptrs, mask=mask, other=0.0)
+
+        Xout = (tl.sum(X, axis=1) + B) / 4
+        tl.store(Xbuf_ptrs, Xout, mask=mask)
+
+except:
+    pb_triton_step_kernel = None
+
+
+class PBTorchCUDAKernelBackend(PBBackend):
+    def __init__(self) -> None:
+        self.kernel = load_inline(
+            name="pb_solver",
+            cpp_sources=[pb_solver_cpp_source],
+            cuda_sources=[pb_solver_cuda_source],
+            functions=["pb_solver_run"],
+            verbose=True,
+        )
+
+    def solve(self, num_iters, A, X, B, Xbuf) -> None:
+        self.kernel.pb_solver_run(A, X, B, Xbuf, num_iters)
+
+
+class PBTritonBackend(PBBackend):
+    def step(
+        self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
+    ) -> None:
+        n_elements = X.shape[0]
+        grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+        assert pb_triton_step_kernel is not None
+        pb_triton_step_kernel[grid](A, X, B, Xbuf, A.stride(0), n_elements, 1024)
+
+    def solve(self, num_iters, A, X, B, Xbuf) -> None:
+        for _ in range(num_iters):
+            self.step(X, A, B, Xbuf)
+            X, Xbuf = Xbuf, X
+
+
+class PBTorchNativeBackend(PBBackend):
+    def step(
+        self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
+    ) -> None:
+        X.copy_(X[A].sum(-1).add_(B).div_(4))
+
+    def solve(self, num_iters, A, X, B, Xbuf) -> None:
+        for _ in range(num_iters):
+            self.step(X, A, B, Xbuf)
+
+
+class PoissonBlendingSolver:
+    def __init__(self, backend: str, device: str):
+        self.backend = backend
+        if backend == "torch-native":
+            self.pb_solver = PBTorchNativeBackend()
+        elif backend == "torch-cuda":
+            self.pb_solver = PBTorchCUDAKernelBackend()
+        elif backend == "triton":
+            self.pb_solver = PBTritonBackend()
+        else:
+            raise ValueError(f"Unknown backend: {backend}")
+
+        self.device = device
+        self.lap_kernel = torch.tensor(
+            [[0, -1, 0], [-1, 4, -1], [0, -1, 0]], device=device, dtype=torch.float32
+        ).view(1, 1, 3, 3)
+        self.lap_kernel4 = torch.tensor(
+            [
+                [[0, -1, 0], [0, 1, 0], [0, 0, 0]],
+                [[0, 0, 0], [0, 1, 0], [0, -1, 0]],
+                [[0, 0, 0], [-1, 1, 0], [0, 0, 0]],
+                [[0, 0, 0], [0, 1, -1], [0, 0, 0]],
+            ],
+            device=device,
+            dtype=torch.float32,
+        ).view(4, 1, 3, 3)
+        self.neighbor_kernel = torch.tensor(
+            [[0, 1, 0], [1, 0, 1], [0, 1, 0]], device=device, dtype=torch.float32
+        ).view(1, 1, 3, 3)
+
+    def __call__(
+        self,
+        src: SINGLE_IMAGE_TYPE,
+        mask: SINGLE_IMAGE_TYPE,
+        tgt: SINGLE_IMAGE_TYPE,
+        num_iters: int,
+        inplace: bool = True,
+        grad_mode: str = "src",
+    ):
+        src = image_to_tensor(src, device=self.device)
+        mask = image_to_tensor(mask, device=self.device)
+        tgt = image_to_tensor(tgt, device=self.device)
+
+        assert src.ndim == 3 and tgt.ndim == 3 and mask.ndim in [2, 3]
+
+        if mask.ndim == 3:
+            mask = mask.mean(-1) > 0.5
+        else:
+            mask = mask > 0.5
+        mask[0, :] = 0
+        mask[-1, :] = 0
+        mask[:, 0] = 0
+        mask[:, -1] = 0
+
+        tgt_masked = torch.where(mask[..., None], torch.zeros_like(tgt), tgt)
+
+        x, y = mask.nonzero(as_tuple=True)
+        N = x.shape[0]
+        index_map = torch.cumsum(mask.reshape(-1).long(), dim=-1).reshape(mask.shape)
+        index_map[~mask] = 0
+
+        if grad_mode == "src":
+            src_lap = F.conv2d(
+                src.permute(2, 0, 1)[:, None],
+                weight=self.lap_kernel,
+                padding=1,
+            )[:, 0].permute(1, 2, 0)
+            lap = src_lap
+        elif grad_mode == "max":
+            src_lap = F.conv2d(
+                src.permute(2, 0, 1)[:, None],
+                weight=self.lap_kernel4,
+                padding=1,
+            )
+            tgt_lap = F.conv2d(
+                tgt.permute(2, 0, 1)[:, None],
+                weight=self.lap_kernel4,
+                padding=1,
+            )
+            lap = (
+                torch.where(src_lap.abs() > tgt_lap.abs(), src_lap, tgt_lap)
+                .sum(1)
+                .permute(1, 2, 0)
+            )
+        elif grad_mode == "avg":
+            src_lap = F.conv2d(
+                src.permute(2, 0, 1)[:, None],
+                weight=self.lap_kernel4,
+                padding=1,
+            )
+            tgt_lap = F.conv2d(
+                tgt.permute(2, 0, 1)[:, None],
+                weight=self.lap_kernel4,
+                padding=1,
+            )
+            lap = (src_lap + tgt_lap).mul(0.5).sum(1).permute(1, 2, 0)
+
+        fq_star = F.conv2d(
+            tgt_masked.permute(2, 0, 1)[:, None],
+            weight=self.neighbor_kernel,
+            padding=1,
+        )[:, 0].permute(1, 2, 0)
+
+        A = torch.zeros(N + 1, 4, device=self.device, dtype=torch.long)
+        X = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
+        B = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
+
+        A[1:] = torch.stack(
+            [
+                index_map[x - 1, y],
+                index_map[x + 1, y],
+                index_map[x, y - 1],
+                index_map[x, y + 1],
+            ],
+            dim=-1,
+        )
+        X[1:] = tgt[x, y]
+        B[1:] = lap[x, y] + fq_star[x, y]
+
+        X_flatten = X.flatten()
+        B_flatten = B.flatten()
+        A_flatten = torch.stack([3 * A, 3 * A + 1, 3 * A + 2], dim=1).reshape(-1, 4)
+
+        buffer = torch.zeros_like(X_flatten)
+
+        self.pb_solver.solve(num_iters, A_flatten, X_flatten, B_flatten, buffer)
+
+        if inplace:
+            tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
+        else:
+            tgt = tgt.clone()
+            tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
+
+        return tgt

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/camera.py

@@ -0,0 +1,223 @@
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+from PIL import Image
+from torch import BoolTensor, FloatTensor
+
+from .utils import LIST_TYPE
+
+
+def list_to_pt(
+    x: LIST_TYPE, dtype: Optional[torch.dtype] = None, device: Optional[str] = None
+) -> torch.Tensor:
+    if isinstance(x, list) or isinstance(x, np.ndarray):
+        return torch.tensor(x, dtype=dtype, device=device)
+    return x.to(dtype=dtype)
+
+
+def get_c2w(
+    elevation_deg: LIST_TYPE,
+    distance: LIST_TYPE,
+    azimuth_deg: Optional[LIST_TYPE],
+    num_views: Optional[int] = 1,
+    device: Optional[str] = None,
+) -> torch.FloatTensor:
+    if azimuth_deg is None:
+        assert (
+            num_views is not None
+        ), "num_views must be provided if azimuth_deg is None."
+        azimuth_deg = torch.linspace(
+            0, 360, num_views + 1, dtype=torch.float32, device=device
+        )[:-1]
+    else:
+        num_views = len(azimuth_deg)
+    azimuth_deg = list_to_pt(azimuth_deg, dtype=torch.float32, device=device)
+    elevation_deg = list_to_pt(elevation_deg, dtype=torch.float32, device=device)
+    camera_distances = list_to_pt(distance, dtype=torch.float32, device=device)
+    elevation = elevation_deg * math.pi / 180
+    azimuth = azimuth_deg * math.pi / 180
+    camera_positions = torch.stack(
+        [
+            camera_distances * torch.cos(elevation) * torch.cos(azimuth),
+            camera_distances * torch.cos(elevation) * torch.sin(azimuth),
+            camera_distances * torch.sin(elevation),
+        ],
+        dim=-1,
+    )
+    center = torch.zeros_like(camera_positions)
+    up = torch.tensor([0, 0, 1], dtype=torch.float32, device=device)[None, :].repeat(
+        num_views, 1
+    )
+    lookat = F.normalize(center - camera_positions, dim=-1)
+    right = F.normalize(torch.cross(lookat, up, dim=-1), dim=-1)
+    up = F.normalize(torch.cross(right, lookat, dim=-1), dim=-1)
+    c2w3x4 = torch.cat(
+        [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
+        dim=-1,
+    )
+    c2w = torch.cat([c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1)
+    c2w[:, 3, 3] = 1.0
+    return c2w
+
+
+def get_projection_matrix(
+    fovy_deg: LIST_TYPE,
+    aspect_wh: float = 1.0,
+    near: float = 0.1,
+    far: float = 100.0,
+    device: Optional[str] = None,
+) -> torch.FloatTensor:
+    fovy_deg = list_to_pt(fovy_deg, dtype=torch.float32, device=device)
+    batch_size = fovy_deg.shape[0]
+    fovy = fovy_deg * math.pi / 180
+    tan_half_fovy = torch.tan(fovy / 2)
+    projection_matrix = torch.zeros(
+        batch_size, 4, 4, dtype=torch.float32, device=device
+    )
+    projection_matrix[:, 0, 0] = 1 / (aspect_wh * tan_half_fovy)
+    projection_matrix[:, 1, 1] = -1 / tan_half_fovy
+    projection_matrix[:, 2, 2] = -(far + near) / (far - near)
+    projection_matrix[:, 2, 3] = -2 * far * near / (far - near)
+    projection_matrix[:, 3, 2] = -1
+    return projection_matrix
+
+
+def get_orthogonal_projection_matrix(
+    batch_size: int,
+    left: float,
+    right: float,
+    bottom: float,
+    top: float,
+    near: float = 0.1,
+    far: float = 100.0,
+    device: Optional[str] = None,
+) -> torch.FloatTensor:
+    projection_matrix = torch.zeros(
+        batch_size, 4, 4, dtype=torch.float32, device=device
+    )
+    projection_matrix[:, 0, 0] = 2 / (right - left)
+    projection_matrix[:, 1, 1] = -2 / (top - bottom)
+    projection_matrix[:, 2, 2] = -2 / (far - near)
+    projection_matrix[:, 0, 3] = -(right + left) / (right - left)
+    projection_matrix[:, 1, 3] = -(top + bottom) / (top - bottom)
+    projection_matrix[:, 2, 3] = -(far + near) / (far - near)
+    projection_matrix[:, 3, 3] = 1
+    return projection_matrix
+
+
+@dataclass
+class Camera:
+    c2w: Optional[torch.FloatTensor]
+    w2c: torch.FloatTensor
+    proj_mtx: torch.FloatTensor
+    mvp_mtx: torch.FloatTensor
+    cam_pos: Optional[torch.FloatTensor]
+
+    def __getitem__(self, index):
+        if isinstance(index, int):
+            sl = slice(index, index + 1)
+        elif isinstance(index, slice):
+            sl = index
+        elif isinstance(index, list):
+            sl = index
+        else:
+            raise NotImplementedError
+
+        return Camera(
+            c2w=self.c2w[sl] if self.c2w is not None else None,
+            w2c=self.w2c[sl],
+            proj_mtx=self.proj_mtx[sl],
+            mvp_mtx=self.mvp_mtx[sl],
+            cam_pos=self.cam_pos[sl] if self.cam_pos is not None else None,
+        )
+
+    def to(self, device: Optional[str] = None):
+        if self.c2w is not None:
+            self.c2w = self.c2w.to(device)
+        self.w2c = self.w2c.to(device)
+        self.proj_mtx = self.proj_mtx.to(device)
+        self.mvp_mtx = self.mvp_mtx.to(device)
+        if self.cam_pos is not None:
+            self.cam_pos = self.cam_pos.to(device)
+
+    def __len__(self):
+        return self.c2w.shape[0]
+
+
+def get_camera(
+    elevation_deg: Optional[LIST_TYPE] = None,
+    distance: Optional[LIST_TYPE] = None,
+    fovy_deg: Optional[LIST_TYPE] = None,
+    azimuth_deg: Optional[LIST_TYPE] = None,
+    num_views: Optional[int] = 1,
+    c2w: Optional[torch.FloatTensor] = None,
+    w2c: Optional[torch.FloatTensor] = None,
+    proj_mtx: Optional[torch.FloatTensor] = None,
+    aspect_wh: float = 1.0,
+    near: float = 0.1,
+    far: float = 100.0,
+    perturb_camera_position: Optional[float] = None,
+    device: Optional[str] = None,
+):
+    if w2c is None:
+        if c2w is None:
+            c2w = get_c2w(elevation_deg, distance, azimuth_deg, num_views, device)
+            if perturb_camera_position is not None:
+                perturbed_pos = (
+                    c2w[:, :3, 3]
+                    + torch.randn_like(c2w[:, :3, 3]) * perturb_camera_position
+                )
+                perturbed_pos = (
+                    F.normalize(perturbed_pos, dim=-1)
+                    * ((c2w[:, :3, 3] ** 2).sum(-1) ** 0.5)[:, None]
+                )
+        camera_positions = c2w[:, :3, 3]
+        w2c = torch.linalg.inv(c2w)
+    else:
+        camera_positions = None
+        c2w = None
+    if proj_mtx is None:
+        proj_mtx = get_projection_matrix(
+            fovy_deg, aspect_wh=aspect_wh, near=near, far=far, device=device
+        )
+    mvp_mtx = proj_mtx @ w2c
+    return Camera(
+        c2w=c2w, w2c=w2c, proj_mtx=proj_mtx, mvp_mtx=mvp_mtx, cam_pos=camera_positions
+    )
+
+
+def get_orthogonal_camera(
+    elevation_deg: LIST_TYPE,
+    distance: LIST_TYPE,
+    left: float,
+    right: float,
+    bottom: float,
+    top: float,
+    azimuth_deg: Optional[LIST_TYPE] = None,
+    num_views: Optional[int] = 1,
+    near: float = 0.1,
+    far: float = 100.0,
+    device: Optional[str] = None,
+):
+    c2w = get_c2w(elevation_deg, distance, azimuth_deg, num_views, device)
+    camera_positions = c2w[:, :3, 3]
+    w2c = torch.linalg.inv(c2w)
+    proj_mtx = get_orthogonal_projection_matrix(
+        batch_size=c2w.shape[0],
+        left=left,
+        right=right,
+        bottom=bottom,
+        top=top,
+        near=near,
+        far=far,
+        device=device,
+    )
+    mvp_mtx = proj_mtx @ w2c
+    return Camera(
+        c2w=c2w, w2c=w2c, proj_mtx=proj_mtx, mvp_mtx=mvp_mtx, cam_pos=camera_positions
+    )

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/cv_ops.py

@@ -0,0 +1,191 @@
+from typing import Optional
+
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import cvcuda
+    torch_to_cvc = lambda x, layout: cvcuda.as_tensor(x, layout)
+
+    cvc_to_torch = lambda x, device: torch.tensor(x.cuda(), device=device)
+except ImportError:
+    cvcuda = None
+    torch_to_cvc = None
+    cvc_to_torch = None
+
+
+def inpaint_cvc(
+    image: torch.Tensor,
+    mask: torch.Tensor,
+    padding_size: int,
+    return_dtype: Optional[torch.dtype] = None,
+):
+    input_dtype = image.dtype
+    input_device = image.device
+
+    image = image.detach()
+    mask = mask.detach()
+
+    if image.dtype != torch.uint8:
+        image = (image * 255).to(torch.uint8)
+    if mask.dtype != torch.uint8:
+        mask = (mask * 255).to(torch.uint8)
+
+    image_cvc = torch_to_cvc(image, "HWC")
+    mask_cvc = torch_to_cvc(mask, "HW")
+    output_cvc = cvcuda.inpaint(image_cvc, mask_cvc, padding_size)
+    output = cvc_to_torch(output_cvc, device=input_device)
+
+    if return_dtype == torch.uint8 or input_dtype == torch.uint8:
+        return output
+    return output.to(dtype=input_dtype) / 255.0
+
+
+def inpaint_torch(
+    image: torch.Tensor,
+    mask: torch.Tensor,
+    padding_size: int,
+    return_dtype: Optional[torch.dtype] = None,
+):
+    """
+    Inpaint masked areas in an image using PyTorch operations.
+
+    Args:
+        image: Input image tensor [H, W, C]
+        mask: Binary mask tensor [H, W] where non-zero values indicate pixels to inpaint
+        padding_size: Size of neighborhood to consider for inpainting
+        return_dtype: Optional dtype for the output tensor
+
+    Returns:
+        Inpainted image tensor
+    """
+    input_dtype = image.dtype
+    input_device = image.device
+
+    image = image.detach()
+    mask = mask.detach()
+
+    if image.dtype != torch.uint8:
+        image = (image * 255).to(torch.uint8)
+    if mask.dtype != torch.uint8:
+        mask = (mask * 255).to(torch.uint8)
+
+    # Convert to float for processing
+    image_float = image.float() / 255.0
+    mask_float = (mask > 0).float()  # 1 for areas to inpaint, 0 for known areas
+
+    # Initialize output with original image
+    output = image_float.clone()
+
+    # Create distance-based weights for neighborhood pixels
+    kernel_size = min(2 * padding_size + 1, 31)  # Limit kernel size to avoid memory issues
+    y, x = torch.meshgrid(
+        torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, device=input_device),
+        torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, device=input_device),
+        indexing="ij"
+    )
+    dist = torch.sqrt(x.float() ** 2 + y.float() ** 2)
+    weights = 1.0 / (dist + 1e-6)
+    weights = weights / weights.sum()
+    weights_kernel = weights.unsqueeze(0).unsqueeze(0)
+
+    # For each channel, compute weighted average of non-masked pixels
+    inpainted = []
+
+    for c in range(image_float.shape[-1]):
+        channel = image_float[..., c]
+
+        # Compute weighted sum of valid pixels
+        valid_values = channel * (1 - mask_float)
+        weighted_sum = F.conv2d(
+            valid_values.unsqueeze(0).unsqueeze(0),
+            weights_kernel,
+            padding=kernel_size // 2
+        ).squeeze(0).squeeze(0)
+
+        # Compute sum of weights for valid pixels
+        valid_weights = F.conv2d(
+            (1 - mask_float).unsqueeze(0).unsqueeze(0),
+            weights_kernel,
+            padding=kernel_size // 2
+        ).squeeze(0).squeeze(0)
+
+        # Avoid division by zero
+        valid_weights = torch.clamp(valid_weights, min=1e-6)
+
+        # Compute weighted average
+        inpainted_channel = weighted_sum / valid_weights
+
+        # Only update masked pixels
+        inpainted_channel = torch.where(mask_float > 0, inpainted_channel, channel)
+        inpainted.append(inpainted_channel)
+
+    # Combine channels
+    output = torch.stack(inpainted, dim=-1)
+
+    # Convert back to appropriate dtype
+    if return_dtype == torch.uint8 or input_dtype == torch.uint8:
+        output = (output * 255).to(torch.uint8)
+    else:
+        output = output.to(dtype=input_dtype)
+
+    return output
+
+
+def batch_inpaint_cvc(
+    images: torch.Tensor,
+    masks: torch.Tensor,
+    padding_size: int,
+    return_dtype: Optional[torch.dtype] = None,
+):
+    output = torch.stack(
+        [
+            inpaint_cvc(image, mask, padding_size, return_dtype)
+            for (image, mask) in zip(images, masks)
+        ],
+        axis=0,
+    )
+    return output
+
+
+def batch_erode(
+    masks: torch.Tensor, kernel_size: int, return_dtype: Optional[torch.dtype] = None
+):
+    input_dtype = masks.dtype
+    input_device = masks.device
+    masks = masks.detach()
+    if masks.dtype != torch.uint8:
+        masks = (masks.float() * 255).to(torch.uint8)
+    masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
+    masks_erode_cvc = cvcuda.morphology(
+        masks_cvc,
+        cvcuda.MorphologyType.ERODE,
+        maskSize=(kernel_size, kernel_size),
+        anchor=(-1, -1),
+    )
+    masks_erode = cvc_to_torch(masks_erode_cvc, device=input_device)[..., 0]
+    if return_dtype == torch.uint8 or input_dtype == torch.uint8:
+        return masks_erode
+    return (masks_erode > 0).to(dtype=input_dtype)
+
+
+def batch_dilate(
+    masks: torch.Tensor, kernel_size: int, return_dtype: Optional[torch.dtype] = None
+):
+    input_dtype = masks.dtype
+    input_device = masks.device
+    masks = masks.detach()
+    if masks.dtype != torch.uint8:
+        masks = (masks.float() * 255).to(torch.uint8)
+    masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
+    masks_dilate_cvc = cvcuda.morphology(
+        masks_cvc,
+        cvcuda.MorphologyType.DILATE,
+        maskSize=(kernel_size, kernel_size),
+        anchor=(-1, -1),
+    )
+    masks_dilate = cvc_to_torch(masks_dilate_cvc, device=input_device)[..., 0]
+    if return_dtype == torch.uint8 or input_dtype == torch.uint8:
+        return masks_dilate
+    return (masks_dilate > 0).to(dtype=input_dtype)

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py

@@ -0,0 +1,522 @@
+import io
+import math
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional, Union
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+from PIL import Image
+from torch import BoolTensor, FloatTensor
+
+from .utils import tensor_to_image
+
+
+def dot(x: torch.FloatTensor, y: torch.FloatTensor) -> torch.FloatTensor:
+    return torch.sum(x * y, -1, keepdim=True)
+
+
+@dataclass
+class TexturedMesh:
+    v_pos: torch.FloatTensor
+    t_pos_idx: torch.LongTensor
+
+    # texture coordinates
+    v_tex: Optional[torch.FloatTensor] = None
+    t_tex_idx: Optional[torch.LongTensor] = None
+
+    # texture map
+    texture: Optional[torch.FloatTensor] = None
+
+    # vertices, faces after vertex merging
+    _stitched_v_pos: Optional[torch.FloatTensor] = None
+    _stitched_t_pos_idx: Optional[torch.LongTensor] = None
+
+    _v_nrm: Optional[torch.FloatTensor] = None
+    _v_tang: Optional[torch.FloatTensor] = None
+
+    @property
+    def v_nrm(self) -> torch.FloatTensor:
+        if self._v_nrm is None:
+            self._v_nrm = self._compute_vertex_normal()
+        return self._v_nrm
+
+    @property
+    def v_tang(self) -> torch.FloatTensor:
+        if self._v_tang is None:
+            self._v_tang = self._compute_tangent()
+        return self._v_tang
+
+    def set_vertex_normal(self, v_nrm: torch.FloatTensor) -> None:
+        assert v_nrm.shape == self.v_pos.shape
+        self._v_nrm = v_nrm.to(self.v_pos)
+
+    def set_stitched_mesh(
+        self, v_pos: torch.FloatTensor, t_pos_idx: torch.LongTensor
+    ) -> None:
+        self._stitched_v_pos = v_pos
+        self._stitched_t_pos_idx = t_pos_idx
+
+    @property
+    def stitched_v_pos(self) -> torch.FloatTensor:
+        if self._stitched_v_pos is None:
+            print("Warning: Stitched vertices not available, using original vertices!")
+            return self.v_pos
+        return self._stitched_v_pos
+
+    @property
+    def stitched_t_pos_idx(self) -> torch.LongTensor:
+        if self._stitched_t_pos_idx is None:
+            print("Warning: Stitched faces not available, using original faces!")
+            return self.t_pos_idx
+        return self._stitched_t_pos_idx
+
+    @property
+    def uv_size(self) -> Optional[int]:
+        if self.texture is None:
+            return None
+        return self.texture.shape[0]
+
+    def _compute_vertex_normal(self) -> torch.FloatTensor:
+        if self._stitched_v_pos is None or self._stitched_t_pos_idx is None:
+            print(
+                "Warning: Stitched vertices and faces not available, computing vertex normals on original mesh, which can be erroneous!"
+            )
+            v_pos, t_pos_idx = self.v_pos, self.t_pos_idx
+        else:
+            v_pos, t_pos_idx = self._stitched_v_pos, self._stitched_t_pos_idx
+
+        i0 = t_pos_idx[:, 0]
+        i1 = t_pos_idx[:, 1]
+        i2 = t_pos_idx[:, 2]
+
+        v0 = v_pos[i0, :]
+        v1 = v_pos[i1, :]
+        v2 = v_pos[i2, :]
+
+        face_normals = torch.cross(v1 - v0, v2 - v0)
+
+        # Splat face normals to vertices
+        v_nrm = torch.zeros_like(v_pos)
+        v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
+        v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
+        v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
+
+        # Normalize, replace zero (degenerated) normals with some default value
+        v_nrm = torch.where(
+            dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
+        )
+        v_nrm = F.normalize(v_nrm, dim=1)
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(v_nrm))
+
+        return v_nrm
+
+    def _compute_tangent(self) -> torch.FloatTensor:
+        vn_idx = [None] * 3
+        pos = [None] * 3
+        tex = [None] * 3
+        for i in range(0, 3):
+            pos[i] = self.v_pos[self.t_pos_idx[:, i]]
+            tex[i] = self.v_tex[self.t_tex_idx[:, i]]
+            vn_idx[i] = self.t_pos_idx[:, i]
+
+        tangents = torch.zeros_like(self._v_nrm)
+        tansum = torch.zeros_like(self._v_nrm)
+
+        # Compute tangent space for each triangle
+        uve1 = tex[1] - tex[0]
+        uve2 = tex[2] - tex[0]
+        pe1 = pos[1] - pos[0]
+        pe2 = pos[2] - pos[0]
+
+        nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2]
+        denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1]
+
+        # Avoid division by zero for degenerated texture coordinates
+        tang = nom / torch.where(
+            denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6)
+        )
+
+        # Update all 3 vertices
+        for i in range(0, 3):
+            idx = vn_idx[i][:, None].repeat(1, 3)
+            tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
+            tansum.scatter_add_(
+                0, idx, torch.ones_like(tang)
+            )  # tansum[n_i] = tansum[n_i] + 1
+        tangents = tangents / tansum
+
+        # Normalize and make sure tangent is perpendicular to normal
+        tangents = F.normalize(tangents, dim=1)
+        tangents = F.normalize(
+            tangents - dot(tangents, self._v_nrm) * self._v_nrm, dim=1
+        )
+
+        v_tng = tangents
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(v_tng))
+
+        return v_tng
+
+    def to(self, device: Optional[str] = None):
+        self.v_pos = self.v_pos.to(device)
+        self.t_pos_idx = self.t_pos_idx.to(device)
+        if self.v_tex is not None:
+            self.v_tex = self.v_tex.to(device)
+        if self.t_tex_idx is not None:
+            self.t_tex_idx = self.t_tex_idx.to(device)
+        if self.texture is not None:
+            self.texture = self.texture.to(device)
+        if self._stitched_v_pos is not None:
+            self._stitched_v_pos = self._stitched_v_pos.to(device)
+        if self._stitched_t_pos_idx is not None:
+            self._stitched_t_pos_idx = self._stitched_t_pos_idx.to(device)
+        if self._v_nrm is not None:
+            self._v_nrm = self._v_nrm.to(device)
+        if self._v_tang is not None:
+            self._v_tang = self._v_tang.to(device)
+
+
+@contextmanager
+def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
+    texture_ = mesh.texture
+    mesh.texture = texture
+    try:
+        yield
+    finally:
+        mesh.texture = texture_
+
+
+def load_mesh(
+    mesh_path: str,
+    rescale: bool = False,
+    move_to_center: bool = False,
+    scale: float = 0.5,
+    flip_uv: bool = True,
+    merge_vertices: bool = True,
+    default_uv_size: Optional[int] = None,
+    shape_init_mesh_up: str = "+y",
+    shape_init_mesh_front: str = "+x",
+    front_x_to_y: bool = False,
+    device: Optional[str] = None,
+    return_transform: bool = False,
+) -> TexturedMesh:
+    if mesh_path.endswith(".npz"):
+
+        class Mesh:
+            vertices = None
+            faces = None
+
+        data = np.load(mesh_path)
+        mesh = Mesh()
+        mesh.vertices = data["vertices"]
+        mesh.faces = data["faces"]
+        merge_vertices = False
+    else:
+        scene = trimesh.load(mesh_path, force="mesh", process=False)
+        if isinstance(scene, trimesh.Trimesh):
+            mesh = scene
+        elif isinstance(scene, trimesh.scene.Scene):
+            mesh = trimesh.Trimesh()
+            for obj in scene.geometry.values():
+                mesh = trimesh.util.concatenate([mesh, obj])
+        else:
+            raise ValueError(f"Unknown mesh type at {mesh_path}.")
+
+    vertex_normals = getattr(mesh, "vertex_normals", None)
+
+    # move to center
+    transform_offset = None
+    if move_to_center:
+        centroid = mesh.vertices.mean(0)
+        mesh.vertices = mesh.vertices - centroid
+        transform_offset = centroid  # record the offset
+
+    # rescale
+    transform_scale = None
+    if rescale:
+        max_scale = np.abs(mesh.vertices).max()
+        mesh.vertices = mesh.vertices / max_scale * scale
+        transform_scale = max_scale / scale
+
+    dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
+    dir2vec = {
+        "+x": np.array([1, 0, 0]),
+        "+y": np.array([0, 1, 0]),
+        "+z": np.array([0, 0, 1]),
+        "-x": np.array([-1, 0, 0]),
+        "-y": np.array([0, -1, 0]),
+        "-z": np.array([0, 0, -1]),
+    }
+    if shape_init_mesh_up not in dirs or shape_init_mesh_front not in dirs:
+        raise ValueError(
+            f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
+        )
+    if shape_init_mesh_up[1] == shape_init_mesh_front[1]:
+        raise ValueError(
+            "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
+        )
+    z_, x_ = (
+        dir2vec[shape_init_mesh_up],
+        dir2vec[shape_init_mesh_front],
+    )
+    y_ = np.cross(z_, x_)
+    std2mesh = np.stack([x_, y_, z_], axis=0).T
+    mesh2std = np.linalg.inv(std2mesh)
+    mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
+    if vertex_normals is not None:
+        vertex_normals = np.dot(mesh2std, vertex_normals.T).T
+    if front_x_to_y:
+        x = mesh.vertices[:, 1].copy()
+        y = -mesh.vertices[:, 0].copy()
+        mesh.vertices[:, 0] = x
+        mesh.vertices[:, 1] = y
+        if vertex_normals is not None:
+            vx = vertex_normals[:, 1].copy()
+            vy = -vertex_normals[:, 0].copy()
+            vertex_normals[:, 0] = vx
+            vertex_normals[:, 1] = vy
+
+    v_pos = torch.tensor(mesh.vertices, dtype=torch.float32)
+    t_pos_idx = torch.tensor(mesh.faces, dtype=torch.int64)
+
+    if hasattr(mesh, "visual") and hasattr(mesh.visual, "uv"):
+        v_tex = torch.tensor(mesh.visual.uv, dtype=torch.float32)
+        if flip_uv:
+            v_tex[:, 1] = 1.0 - v_tex[:, 1]
+        t_tex_idx = t_pos_idx.clone()
+        if (
+            default_uv_size is not None
+            or getattr(mesh.visual.material, "baseColorTexture", None) is None
+        ):
+            assert default_uv_size is not None
+            texture = torch.zeros(
+                (default_uv_size, default_uv_size, 3), dtype=torch.float32
+            )
+        else:
+            texture = torch.tensor(
+                np.array(mesh.visual.material.baseColorTexture) / 255.0,
+                dtype=torch.float32,
+            )[..., :3]
+    else:
+        v_tex = None
+        t_tex_idx = None
+        texture = None
+
+    textured_mesh = TexturedMesh(
+        v_pos=v_pos,
+        t_pos_idx=t_pos_idx,
+        v_tex=v_tex,
+        t_tex_idx=t_tex_idx,
+        texture=texture,
+    )
+
+    if vertex_normals is not None:
+        v_nrm = F.normalize(torch.tensor(vertex_normals, dtype=torch.float32), dim=-1)
+        textured_mesh.set_vertex_normal(v_nrm)
+
+    if merge_vertices and vertex_normals is None:
+        # only merge vertices when vertex normals are not available
+        mesh.merge_vertices(merge_tex=True)
+        textured_mesh.set_stitched_mesh(
+            torch.tensor(mesh.vertices, dtype=torch.float32),
+            torch.tensor(mesh.faces, dtype=torch.int64),
+        )
+    else:
+        textured_mesh.set_stitched_mesh(textured_mesh.v_pos, textured_mesh.t_pos_idx)
+
+    textured_mesh.to(device)
+
+    if return_transform:
+        return textured_mesh, transform_offset, transform_scale
+    else:
+        return textured_mesh
+
+
+def replace_mesh_texture_and_save_trimesh(
+    input_path: str,
+    output_path: str,
+    texture: Union[Image.Image, np.ndarray, torch.Tensor],
+    metallic_roughness_texture: Optional[
+        Union[Image.Image, np.ndarray, torch.Tensor]
+    ] = None,
+    normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
+    texture_format: str = "JPEG",
+    task_id: str = "",
+    **kwargs,
+) -> None:
+    tmesh = trimesh.load(input_path, force="mesh", process=False)
+
+    def convert_image(tensor):
+        texture = tensor_to_image(tensor)
+        texture.format = texture_format
+        return texture
+
+    texture = convert_image(texture)
+    if metallic_roughness_texture is not None:
+        metallic_roughness_texture = convert_image(metallic_roughness_texture)
+
+    if normal_texture is not None:
+        normal_texture = convert_image(normal_texture)
+
+    if not isinstance(tmesh.visual.material, trimesh.visual.material.PBRMaterial):
+        material = trimesh.visual.texture.PBRMaterial(
+            baseColorFactor=[1.0, 1.0, 1.0, 1.0],
+            roughnessFactor=0.9 if metallic_roughness_texture is None else 1.0,
+            metallicFactor=0.0 if metallic_roughness_texture is None else 1.0,
+            baseColorTexture=texture,
+            metallicRoughnessTexture=metallic_roughness_texture,
+            normalTexture=normal_texture,
+        )
+        tmesh.visual = trimesh.visual.TextureVisuals(
+            uv=tmesh.visual.uv, image=texture, material=material
+        )
+    else:
+        tmesh.visual.material.baseColorTexture = texture
+        tmesh.visual.material.baseColorFactor = [1.0, 1.0, 1.0, 1.0]
+        if metallic_roughness_texture is None:
+            tmesh.visual.material.roughnessFactor = 0.9
+            tmesh.visual.material.metallicFactor = 0.0
+        else:
+            tmesh.visual.material.roughnessFactor = 1.0
+            tmesh.visual.material.metallicFactor = 1.0
+            tmesh.visual.material.metallicRoughnessTexture = metallic_roughness_texture
+            tmesh.visual.material.normalTexture = normal_texture
+
+    def p_func(tree, task_id_in=task_id):
+        if "nodes" in tree and tree["nodes"][0]["name"].find("world") >= 0:
+            tree["nodes"].pop(0)
+        tree["asset"]["generator"] = "https://github.com/huanngzh/MV-Adapter"
+        change_dict = {
+            "nodes": "node",
+            "meshes": "mesh",
+            "materials": "mat",
+            "images": "image",
+        }
+        for k, v in change_dict.items():
+            if k in tree:
+                for index in range(len(tree[k])):
+                    tree[k][index]["name"] = f"mvadapter_{v}_{task_id_in}"
+
+    tmesh.export(output_path, tree_postprocessor=p_func)
+
+
+# def replace_mesh_texture_and_save_gltflib(
+#     input_path: str,
+#     output_path: str,
+#     texture: Union[Image.Image, np.ndarray, torch.Tensor],
+#     metallic_roughness_texture: Optional[
+#         Union[Image.Image, np.ndarray, torch.Tensor]
+#     ] = None,
+#     normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
+#     texture_format: str = "JPEG",
+#     normal_strength: float = 1.0,
+#     task_id: str = "",
+#     **kwargs,
+# ) -> None:
+#     from gltflib import GLTF
+#     from gltflib import Image as gltfImage
+#     from gltflib import NormalTextureInfo, Texture, TextureInfo
+
+#     def convert_image(tensor):
+#         texture = tensor_to_image(tensor)
+#         texture.format = texture_format
+#         return texture
+
+#     def add_texture(gltf: GLTF, image: Image.Image) -> int:
+#         jpeg_buffer = io.BytesIO()
+#         image.save(jpeg_buffer, format=texture_format)
+#         new_texture_data = bytearray(jpeg_buffer.getvalue())
+#         _, offset, bytelen = gltf._create_or_extend_glb_resource(new_texture_data)
+#         bufferView = gltf._create_embedded_image_buffer_view(offset, bytelen)
+#         image = gltfImage(mimeType="image/jpeg", bufferView=bufferView)
+#         if gltf.model.images:
+#             gltf.model.images.append(image)
+#         else:
+#             gltf.model.images = [image]
+#         texture = Texture(source=len(gltf.model.images) - 1)
+#         if gltf.model.textures:
+#             gltf.model.textures.append(texture)
+#         else:
+#             gltf.model.textures = [texture]
+#         return len(gltf.model.textures) - 1
+
+#     gltf = GLTF.load(input_path)
+
+#     texture = convert_image(texture)
+#     if metallic_roughness_texture is not None:
+#         metallic_roughness_texture = convert_image(metallic_roughness_texture)
+
+#     if normal_texture is not None:
+#         normal_texture = convert_image(normal_texture)
+
+#     pbr_material = gltf.model.materials[0].pbrMetallicRoughness
+
+#     pbr_material.baseColorTexture = TextureInfo(index=add_texture(gltf, texture))
+#     pbr_material.baseColorFactor = [1, 1, 1, 1]
+
+#     if metallic_roughness_texture is not None:
+#         pbr_material.metallicRoughnessTexture = TextureInfo(
+#             index=add_texture(gltf, metallic_roughness_texture)
+#         )
+#         pbr_material.metallicFactor = 1.0
+#         pbr_material.roughnessFactor = 1.0
+#     else:
+#         pbr_material.metallicFactor = 0.0
+#         pbr_material.roughnessFactor = 0.9
+
+#     if normal_texture is not None:
+#         gltf.model.materials[0].normalTexture = NormalTextureInfo(
+#             index=add_texture(gltf, normal_texture), scale=normal_strength
+#         )
+
+#     gltf.model.nodes[0].name = f"mvadapter_node_{task_id}"
+#     gltf.model.meshes[0].name = f"mvadapter_mesh_{task_id}"
+#     gltf.model.materials[0].name = f"mvadapter_material_{task_id}"
+#     gltf.model.asset.generator = "https://github.com/huanngzh/MV-Adapter"
+#     gltf.export(output_path)
+
+
+def replace_mesh_texture_and_save(
+    input_path: str,
+    output_path: str,
+    texture: Union[Image.Image, np.ndarray, torch.Tensor],
+    metallic_roughness_texture: Optional[
+        Union[Image.Image, np.ndarray, torch.Tensor]
+    ] = None,
+    normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
+    normal_strength: float = 1.0,
+    texture_format: str = "JPEG",
+    task_id: str = "",
+    backend: str = "trimesh",
+) -> None:
+    if backend == "trimesh":
+        replace_mesh_texture_and_save_trimesh(
+            input_path,
+            output_path,
+            texture,
+            metallic_roughness_texture=metallic_roughness_texture,
+            normal_texture=normal_texture,
+            texture_format=texture_format,
+            task_id=task_id,
+        )
+    # elif backend == "gltflib":
+    #     replace_mesh_texture_and_save_gltflib(
+    #         input_path,
+    #         output_path,
+    #         texture,
+    #         metallic_roughness_texture=metallic_roughness_texture,
+    #         normal_texture=normal_texture,
+    #         normal_strength=normal_strength,
+    #         task_id=task_id,
+    #     )
+    else:
+        raise NotImplementedError

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh_process.py

@@ -0,0 +1,348 @@
+import numpy as np
+import open3d as o3d
+import pymeshlab
+import torch
+import trimesh
+from pymeshlab import PercentageValue
+
+
+### Mesh Utils ###
+##### read mesh
+def read_mesh_from_path(mesh_path):
+    ms = pymeshlab.MeshSet()
+    ms.load_new_mesh(mesh_path)
+    return ms
+
+
+def mesh_to_meshlab(vertices, faces):
+    mesh = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
+    ms = pymeshlab.MeshSet()
+    ms.add_mesh(mesh)
+    return ms
+
+
+def meshlab_to_mesh(ms):
+    m = ms.current_mesh()
+    return m.vertex_matrix(), m.face_matrix(), m.vertex_normal_matrix()
+
+
+##### decimation
+def decimate_quadric_edge_collapse_with_texture(
+    ms, targetfacenum=None, preservenormal=True, verbose=False
+):
+    # targetfacenum: int, Target number of faces.
+    # preservenormal: bool, Preserve the normals of the original mesh.
+    if verbose:
+        print("Starting decimation ...  ")
+    m = ms.current_mesh()
+    if targetfacenum is None:
+        targetfacenum = int(m.face_number() * 0.5)
+    if verbose:
+        print("... Initial face number is %d ... " % m.face_number())
+    ms.meshing_decimation_quadric_edge_collapse_with_texture(
+        targetfacenum=targetfacenum, preservenormal=preservenormal
+    )
+    if verbose:
+        print("... Decimated face number is %d ... " % m.face_number())
+        print("Decimation done!\n ")
+
+
+def decimate_quadric_edge_collapse(
+    ms, targetfacenum=None, preservenormal=True, verbose=False
+):
+    # targetfacenum: int, Target number of faces.
+    # preservenormal: bool, Preserve the normals of the original mesh.
+    if verbose:
+        print("Starting decimation ...  ")
+    m = ms.current_mesh()
+    if targetfacenum is None:
+        targetfacenum = int(m.face_number() * 0.5)
+    if verbose:
+        print("... Initial face number is %d ... " % m.face_number())
+    ms.meshing_decimation_quadric_edge_collapse(
+        targetfacenum=targetfacenum, preservenormal=preservenormal
+    )
+    if verbose:
+        print("... Decimated face number is %d ... " % m.face_number())
+        print("Decimation done!\n ")
+
+
+##### vertex merge
+def merge_close_vertices(ms, threshold=0.0001, verbose=False):
+    # threshold: float, Merge together all the vertices that are nearer than the specified threshold.
+    if verbose:
+        print("Starting merge vertices ...  ")
+    m = ms.current_mesh()
+    if verbose:
+        print("... Initial vertex number is %d ... " % m.vertex_number())
+    ms.meshing_merge_close_vertices(threshold=PercentageValue(threshold * 100))
+    if verbose:
+        print("... Merged vertex number is %d ... " % m.vertex_number())
+        print("Merge vertices done!\n ")
+
+
+##### Island Removal
+def remove_isolated_pieces(ms, mincomponentsize=25, diameter=None, verbose=False):
+    # mincomponentsize: Delete isolated connected components composed by a limited number of triangles
+    # diameter: Delete isolated connected components whose diameter is smaller than the specified constant
+    if verbose:
+        print("Starting remove isolated pieces ...  ")
+    m = ms.current_mesh()
+    if verbose:
+        print("... Initial face number is %d ... " % m.face_number())
+    if diameter is None:
+        ms.meshing_remove_connected_component_by_face_number(
+            mincomponentsize=mincomponentsize, removeunref=True
+        )
+    else:
+        ms.meshing_remove_connected_component_by_diameter(
+            mincomponentdiag=PercentageValue(diameter), removeunref=True
+        )
+    if verbose:
+        print("... Isolated removed face number is %d ... " % m.face_number())
+        print("Remove isolated pieces done!\n ")
+
+
+##### hole filling
+def fix_hole(ms, maxholesize=30, verbose=False):
+    # maxholesize: int, Maximum size of the hole to be filled.
+    if verbose:
+        print("Starting fix holes ...  ")
+    m = ms.current_mesh()
+    if verbose:
+        print("... Initial face number is %d ... " % m.face_number())
+    ms.meshing_close_holes(maxholesize=maxholesize)
+    if verbose:
+        print("... Fixed hole face number is %d ... " % m.face_number())
+        print("Fix holes done!\n ")
+
+
+##### repair non manifold edges
+def repair_non_manifold(ms, verbose=False):
+    if verbose:
+        print("Starting repair non manifold edges ...  ")
+    m = ms.current_mesh()
+    if verbose:
+        print("... Initial face number is %d ... " % m.face_number())
+    ms.meshing_repair_non_manifold_edges()
+    ms.meshing_repair_non_manifold_vertices(vertdispratio=0.1)
+    ms.meshing_remove_duplicate_faces()
+    if verbose:
+        print("... Fixed non manifold edges face number is %d ... " % m.face_number())
+        print("Repair non manifold edges done!\n ")
+
+
+##### laplacian_smooth
+def laplacian_smooth(ms, stepsmoothnum=3, verbose=False):
+    # stepsmoothnum: int, Number of smoothing steps to be performed
+    if verbose:
+        print("Starting laplacian smooth ...  ")
+    m = ms.current_mesh()
+    ms.apply_coord_laplacian_smoothing(stepsmoothnum=stepsmoothnum)
+    if verbose:
+        print("Laplacian smooth done!\n ")
+
+
+##### taubin_smooth
+def taubin_smooth(ms, stepsmoothnum=3, verbose=False):
+    if verbose:
+        print("Starting Taubin smooth ...  ")
+    m = ms.current_mesh()
+    ms.apply_coord_taubin_smoothing(stepsmoothnum=stepsmoothnum)
+    if verbose:
+        print("Taubin smooth done!\n ")
+
+
+##### compute_normal
+def compute_normal(ms, weightmode="Simple Average", verbose=False):
+    if verbose:
+        print("Starting compute_normal_per_vertex ...  ")
+    m = ms.current_mesh()
+    ms.compute_normal_per_vertex(weightmode=weightmode)
+    if verbose:
+        print("compute_normal_per_vertex done!\n ")
+
+
+### Pre-process Mesh ###
+def process_mesh(
+    vertices,
+    faces,
+    threshold=0.0001,
+    mincomponentRatio=0.02,
+    targetfacenum=50000,
+    maxholesize=30,
+    stepsmoothnum=10,
+    verbose=False,
+):
+    ms = mesh_to_meshlab(vertices, faces)
+
+    ### Vertex Merge
+    merge_close_vertices(ms, threshold=threshold, verbose=verbose)
+
+    ### Island Removal
+    faces = ms.current_mesh().face_matrix()
+    remove_isolated_pieces(
+        ms, mincomponentsize=int(len(faces) * mincomponentRatio), verbose=verbose
+    )
+
+    ### Hole Filling
+    repair_non_manifold(ms)  # repair before fix hole
+    fix_hole(ms, maxholesize=maxholesize, verbose=verbose)
+
+    ### Taubin Smoothing
+    taubin_smooth(ms, stepsmoothnum=stepsmoothnum, verbose=verbose)
+
+    vertices, faces, _ = meshlab_to_mesh(ms)
+    if faces.shape[0] > targetfacenum:
+        device = o3d.core.Device("CPU:0")
+        dtype_f = o3d.core.float32
+        dtype_i = o3d.core.int64
+        mesh = o3d.t.geometry.TriangleMesh(device)
+        mesh.vertex.positions = o3d.core.Tensor(
+            vertices.astype(np.float32), dtype_f, device
+        )
+        mesh.triangle.indices = o3d.core.Tensor(faces.astype(np.int64), dtype_i, device)
+        simplified_mesh = mesh.simplify_quadric_decimation(
+            target_reduction=1.0 - float(targetfacenum) / faces.shape[0]
+        )
+        ms.clear()
+        vertices = simplified_mesh.vertex.positions.numpy()
+        faces = simplified_mesh.triangle.indices.numpy()
+        mesh = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
+        ms.add_mesh(mesh)
+
+    ### Mesh Simplification/Decimation
+    # decimate_quadric_edge_collapse(ms, targetfacenum=targetfacenum, verbose=verbose)
+    taubin_smooth(ms, stepsmoothnum=stepsmoothnum, verbose=verbose)
+    repair_non_manifold(ms, verbose=verbose)
+    compute_normal(ms, verbose=verbose)
+    return meshlab_to_mesh(ms)
+
+
+### UV Un-Warp ###
+def uv_parameterize_uvatlas(
+    vertices,
+    faces,
+    size=1024,
+    gutter=2.5,
+    max_stretch=0.1666666716337204,
+    parallel_partitions=16,
+    nthreads=0,
+):
+    device = o3d.core.Device("CPU:0")
+    dtype_f = o3d.core.float32
+    dtype_i = o3d.core.int64
+
+    mesh = o3d.t.geometry.TriangleMesh(device)
+
+    mesh.vertex.positions = o3d.core.Tensor(
+        vertices.astype(np.float32), dtype_f, device
+    )
+    mesh.triangle.indices = o3d.core.Tensor(faces.astype(np.int64), dtype_i, device)
+
+    mesh.compute_uvatlas(
+        size=size,
+        gutter=gutter,
+        max_stretch=max_stretch,
+        parallel_partitions=parallel_partitions,
+        nthreads=nthreads,
+    )
+
+    return mesh.triangle.texture_uvs.numpy()  # (#F, 3, 2)
+
+
+### Pack All ###
+def process_raw(mesh_path, save_path, preprocess=True, device="cpu"):
+    scene = trimesh.load(mesh_path, force="mesh", process=False)
+    if isinstance(scene, trimesh.Trimesh):
+        mesh = scene
+    elif isinstance(scene, trimesh.scene.Scene):
+        mesh = trimesh.Trimesh()
+        for obj in scene.geometry.values():
+            mesh = trimesh.util.concatenate([mesh, obj])
+    else:
+        raise ValueError(f"Unknown mesh type at {mesh_path}.")
+
+    vertices = mesh.vertices
+    faces = mesh.faces
+
+    mesh_post_process_options = {
+        "mincomponentRatio": 0.02,
+        "targetfacenum": 50000,
+        "maxholesize": 100,
+        "stepsmoothnum": 10,
+        "verbose": False,
+    }
+
+    if preprocess:
+        v_pos, t_pos_idx, normals = process_mesh(
+            vertices=vertices,
+            faces=faces,
+            **mesh_post_process_options,
+        )
+    else:
+        v_pos, t_pos_idx, normals = vertices, faces, mesh.vertex_normals
+
+    v_tex_np = (
+        uv_parameterize_uvatlas(v_pos, t_pos_idx).reshape(-1, 2).astype(np.float32)
+    )
+
+    v_pos = torch.from_numpy(v_pos).to(device=device, dtype=torch.float32)
+    t_pos_idx = torch.from_numpy(t_pos_idx).to(device=device, dtype=torch.long)
+    v_tex = torch.from_numpy(v_tex_np).to(device=device, dtype=torch.float32)
+    normals = torch.from_numpy(normals).to(device=device, dtype=torch.float32)
+    assert v_tex.shape[0] == t_pos_idx.shape[0] * 3
+    t_tex_idx = torch.arange(
+        t_pos_idx.shape[0] * 3,
+        device=device,
+        dtype=torch.long,
+    ).reshape(-1, 3)
+    # uv, index = torch.unique(v_tex, dim=0, return_inverse=True) # 这样实现是2毫秒
+    # super efficient de-duplication
+    v_tex_u_uint32 = v_tex_np[..., 0].view(np.uint32)
+    v_tex_v_uint32 = v_tex_np[..., 1].view(np.uint32)
+    v_hashed = (v_tex_u_uint32.astype(np.uint64) << 32) | v_tex_v_uint32
+    v_hashed = torch.from_numpy(v_hashed.view(np.int64)).to(v_pos.device)
+
+    t_pos_idx_f3 = torch.arange(
+        t_pos_idx.shape[0] * 3, device=t_pos_idx.device, dtype=torch.long
+    ).reshape(-1, 3)
+    v_pos_f3 = v_pos[t_pos_idx].reshape(-1, 3)
+    normals_f3 = normals[t_pos_idx].reshape(-1, 3)
+
+    v_hashed_dedup, inverse_indices = torch.unique(v_hashed, return_inverse=True)
+    dedup_size, full_size = v_hashed_dedup.shape[0], inverse_indices.shape[0]
+    indices = torch.scatter_reduce(
+        torch.full(
+            [dedup_size],
+            fill_value=full_size,
+            device=inverse_indices.device,
+            dtype=torch.long,
+        ),
+        index=inverse_indices,
+        src=torch.arange(full_size, device=inverse_indices.device, dtype=torch.int64),
+        dim=0,
+        reduce="amin",
+    )
+    v_tex = v_tex[indices]
+    t_tex_idx = inverse_indices.reshape(-1, 3)
+
+    v_pos = v_pos_f3[indices]
+    normals = normals_f3[indices]
+
+    normals = normals.to(dtype=torch.float32, device=device)
+
+    # either flip uv or flip texture
+    # here we flip uv
+    uv_to_save = v_tex.clone()
+    uv_to_save[:, 1] = 1.0 - uv_to_save[:, 1]
+
+    visual = trimesh.visual.TextureVisuals(uv=uv_to_save.cpu().numpy())
+    tmesh = trimesh.Trimesh(
+        vertices=v_pos.cpu().numpy(),
+        faces=t_tex_idx.cpu().numpy(),
+        vertex_normals=normals.cpu().numpy(),
+        visual=visual,
+        process=False,
+    )
+    tmesh.export(save_path)

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/projection.py

@@ -0,0 +1,204 @@
+import glob
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from einops import rearrange
+from PIL import Image
+
+from .blend import PoissonBlendingSolver
+from .camera import Camera, get_camera, get_orthogonal_camera
+from .mesh import TexturedMesh, load_mesh, replace_mesh_texture_and_save
+from .render import NVDiffRastContextWrapper, SimpleNormalization, render
+from .seg import SegmentationModel
+from .utils import (
+    IMAGE_TYPE,
+    LIST_TYPE,
+    image_to_tensor,
+    make_image_grid,
+    tensor_to_image,
+)
+from .uv import (
+    ExponentialBlend,
+    SimpleUVValidityStrategy,
+    uv_blend,
+    uv_precompute,
+    uv_render_attr,
+    uv_render_geometry,
+)
+from .warp import compute_warp_field
+
+
+@dataclass
+class CameraProjectionOutput:
+    uv_proj: torch.FloatTensor
+    uv_proj_mask: torch.BoolTensor
+    uv_depth_grad: Optional[torch.FloatTensor]
+    uv_aoi_cos: Optional[torch.FloatTensor]
+
+
+class CameraProjection:
+    def __init__(
+        self,
+        pb_backend: str,
+        bg_remover: Optional[SegmentationModel],
+        device: str,
+        context_type: str = "gl",
+    ) -> None:
+        self.pb_solver = PoissonBlendingSolver(pb_backend, device)
+        self.ctx = NVDiffRastContextWrapper(device, context_type)
+        self.bg_remover = bg_remover
+        self.device = device
+
+    def __call__(
+        self,
+        images: IMAGE_TYPE,
+        mesh: TexturedMesh,
+        cam: Optional[Camera] = None,
+        fovy_deg: Optional[LIST_TYPE] = None,
+        masks: Optional[IMAGE_TYPE] = None,
+        remove_bg: bool = False,
+        c2w: Optional[torch.FloatTensor] = None,
+        elevation_deg: Optional[LIST_TYPE] = None,
+        distance: Optional[LIST_TYPE] = None,
+        azimuth_deg: Optional[LIST_TYPE] = None,
+        num_views: Optional[int] = None,
+        uv_size: int = 2048,
+        warp_images: bool = False,
+        images_background: Optional[float] = None,
+        iou_rejection_threshold: Optional[float] = 0.8,
+        aoi_cos_valid_threshold: float = 0.3,
+        depth_grad_dilation: int = 5,
+        depth_grad_threshold: float = 0.1,
+        uv_exp_blend_alpha: float = 6,
+        uv_exp_blend_view_weight: Optional[torch.FloatTensor] = None,
+        poisson_blending: bool = True,
+        pb_num_iters: int = 1000,
+        pb_keep_original_border: bool = True,
+        from_scratch: bool = False,
+        uv_padding: bool = True,
+        return_uv_projection_mask: bool = False,
+        return_dict: bool = False,
+    ) -> Optional[torch.FloatTensor]:
+        images_pt = image_to_tensor(images, device=self.device)
+        assert images_pt.ndim == 4
+        Nv, H, W, _ = images_pt.shape
+
+        if masks is not None:
+            masks_pt = image_to_tensor(masks, device=self.device)
+        else:
+            if remove_bg:
+                assert self.bg_remover is not None
+                masks_pt = self.bg_remover(images_pt)
+            else:
+                masks_pt = None
+
+        if masks_pt is not None and masks_pt.ndim == 4:
+            masks_pt = masks_pt.mean(-1)
+
+        if cam is None:
+            cam = get_camera(
+                elevation_deg,
+                distance,
+                fovy_deg,
+                azimuth_deg,
+                num_views,
+                c2w,
+                aspect_wh=W / H,
+                device=self.device,
+            )
+        uv_precompute_output = uv_precompute(
+            self.ctx, mesh, height=uv_size, width=uv_size
+        )
+        uv_render_geometry_output = uv_render_geometry(
+            self.ctx,
+            mesh,
+            cam,
+            view_height=H,
+            view_width=W,
+            uv_precompute_output=uv_precompute_output,
+            compute_depth_grad=True,
+            depth_grad_dilation=depth_grad_dilation,
+        )
+
+        # IoU rejection
+        if masks_pt is not None and iou_rejection_threshold is not None:
+            given_masks = (masks_pt > 0.5).float()
+            render_masks = uv_render_geometry_output.view_mask.float()
+            intersection = given_masks * render_masks
+            union = given_masks + render_masks - intersection
+            iou = intersection.sum((1, 2)) / union.sum((1, 2))
+            iou_min = iou.min()
+            print(f"Debug: Per view IoU: {iou.tolist()}")
+            if iou_min < iou_rejection_threshold:
+                print(
+                    f"Warning: Minimum view IoU {iou_min} below threshold {iou_rejection_threshold}, skipping camera projection!"
+                )
+                return None
+
+        if warp_images:
+            # TODO: clean code
+            assert images_background is not None
+            render_attr = render(
+                self.ctx,
+                mesh,
+                cam,
+                height=H,
+                width=W,
+                render_attr=True,
+                attr_background=images_background,
+            ).attr
+            images_pt = compute_warp_field(
+                self.ctx.ctx,
+                images_pt,
+                render_attr,
+                n_grid=10,
+                optim_res=[64, 128],
+                optim_step_per_res=20,
+                lambda_reg=2.0,
+                temp_dir="debug_warp",
+                verbose=False,
+                device=self.device,
+            )
+
+        uv_render_attr_output = uv_render_attr(
+            images=images_pt,
+            masks=masks_pt,
+            uv_render_geometry_output=uv_render_geometry_output,
+        )
+        uv_blend_output = uv_blend(
+            uv_precompute_output,
+            uv_render_geometry_output,
+            uv_render_attr_output,
+            uv_validity_strategy=SimpleUVValidityStrategy(
+                aoi_cos_thresh=aoi_cos_valid_threshold,
+                depth_grad_thresh=depth_grad_threshold,
+            ),
+            uv_blend_weight_strategy=ExponentialBlend(
+                alpha=uv_exp_blend_alpha, view_weight=uv_exp_blend_view_weight
+            ),
+            empty_value=1.0,
+            do_uv_padding=uv_padding,
+            pad_unseen_area=from_scratch,
+            poisson_blending=poisson_blending,
+            pb_solver=self.pb_solver,
+            pb_num_iters=pb_num_iters,
+            pb_keep_original_border=pb_keep_original_border,
+        )
+
+        if return_dict:
+            # recommonded new way to get return value
+            return CameraProjectionOutput(
+                uv_proj=uv_blend_output.uv_attr_blend,
+                uv_proj_mask=uv_blend_output.uv_valid_mask_blend,
+                uv_depth_grad=uv_render_geometry_output.uv_depth_grad,
+                uv_aoi_cos=uv_render_geometry_output.uv_aoi_cos,
+            )
+        else:
+            if return_uv_projection_mask:
+                return (
+                    uv_blend_output.uv_attr_blend,
+                    uv_blend_output.uv_valid_mask_blend,
+                )
+            return uv_blend_output.uv_attr_blend

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/render.py

@@ -0,0 +1,310 @@
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional, Union
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+from PIL import Image
+from torch import BoolTensor, FloatTensor
+
+from .camera import Camera
+from .mesh import TexturedMesh
+from .utils import get_clip_space_position, transform_points_homo
+
+
+@dataclass
+class RenderOutput:
+    attr: Optional[torch.FloatTensor] = None
+    mask: Optional[torch.BoolTensor] = None
+    depth: Optional[torch.FloatTensor] = None
+    normal: Optional[torch.FloatTensor] = None
+    tangent: Optional[torch.FloatTensor] = None
+    pos: Optional[torch.FloatTensor] = None
+
+
+class NVDiffRastContextWrapper:
+    def __init__(self, device: str, context_type: str = "gl"):
+        if context_type == "gl":
+            self.ctx = dr.RasterizeGLContext(device=device)
+        elif context_type == "cuda":
+            self.ctx = dr.RasterizeCudaContext(device=device)
+        else:
+            raise NotImplementedError
+
+    def rasterize(self, pos, tri, resolution, ranges=None, grad_db=True):
+        """
+        Rasterize triangles.
+
+        All input tensors must be contiguous and reside in GPU memory except for the ranges tensor that, if specified, has to reside in CPU memory. The output tensors will be contiguous and reside in GPU memory.
+
+        Arguments:
+        glctx	Rasterizer context of type RasterizeGLContext or RasterizeCudaContext.
+        pos	Vertex position tensor with dtype torch.float32. To enable range mode, this tensor should have a 2D shape [num_vertices, 4]. To enable instanced mode, use a 3D shape [minibatch_size, num_vertices, 4].
+        tri	Triangle tensor with shape [num_triangles, 3] and dtype torch.int32.
+        resolution	Output resolution as integer tuple (height, width).
+        ranges	In range mode, tensor with shape [minibatch_size, 2] and dtype torch.int32, specifying start indices and counts into tri. Ignored in instanced mode.
+        grad_db	Propagate gradients of image-space derivatives of barycentrics into pos in backward pass. Ignored if using an OpenGL context that was not configured to output image-space derivatives.
+        Returns:
+        A tuple of two tensors. The first output tensor has shape [minibatch_size, height, width, 4] and contains the main rasterizer output in order (u, v, z/w, triangle_id). If the OpenGL context was configured to output image-space derivatives of barycentrics, the second output tensor will also have shape [minibatch_size, height, width, 4] and contain said derivatives in order (du/dX, du/dY, dv/dX, dv/dY). Otherwise it will be an empty tensor with shape [minibatch_size, height, width, 0].
+        """
+        return dr.rasterize(
+            self.ctx,
+            pos.float().contiguous(),
+            tri.int().contiguous(),
+            resolution,
+            ranges,
+            grad_db,
+        )
+
+    def interpolate(self, attr, rast, tri, rast_db=None, diff_attrs=None):
+        """
+        Interpolate vertex attributes.
+
+        All input tensors must be contiguous and reside in GPU memory. The output tensors will be contiguous and reside in GPU memory.
+
+        Arguments:
+        attr	Attribute tensor with dtype torch.float32. Shape is [num_vertices, num_attributes] in range mode, or [minibatch_size, num_vertices, num_attributes] in instanced mode. Broadcasting is supported along the minibatch axis.
+        rast	Main output tensor from rasterize().
+        tri	Triangle tensor with shape [num_triangles, 3] and dtype torch.int32.
+        rast_db	(Optional) Tensor containing image-space derivatives of barycentrics, i.e., the second output tensor from rasterize(). Enables computing image-space derivatives of attributes.
+        diff_attrs	(Optional) List of attribute indices for which image-space derivatives are to be computed. Special value 'all' is equivalent to list [0, 1, ..., num_attributes - 1].
+        Returns:
+        A tuple of two tensors. The first output tensor contains interpolated attributes and has shape [minibatch_size, height, width, num_attributes]. If rast_db and diff_attrs were specified, the second output tensor contains the image-space derivatives of the selected attributes and has shape [minibatch_size, height, width, 2 * len(diff_attrs)]. The derivatives of the first selected attribute A will be on channels 0 and 1 as (dA/dX, dA/dY), etc. Otherwise, the second output tensor will be an empty tensor with shape [minibatch_size, height, width, 0].
+        """
+        return dr.interpolate(
+            attr.float().contiguous(), rast, tri.int().contiguous(), rast_db, diff_attrs
+        )
+
+    def texture(
+        self,
+        tex,
+        uv,
+        uv_da=None,
+        mip_level_bias=None,
+        mip=None,
+        filter_mode="auto",
+        boundary_mode="wrap",
+        max_mip_level=None,
+    ):
+        """
+        Perform texture sampling.
+
+        All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory.
+
+        Arguments:
+        tex	Texture tensor with dtype torch.float32. For 2D textures, must have shape [minibatch_size, tex_height, tex_width, tex_channels]. For cube map textures, must have shape [minibatch_size, 6, tex_height, tex_width, tex_channels] where tex_width and tex_height are equal. Note that boundary_mode must also be set to 'cube' to enable cube map mode. Broadcasting is supported along the minibatch axis.
+        uv	Tensor containing per-pixel texture coordinates. When sampling a 2D texture, must have shape [minibatch_size, height, width, 2]. When sampling a cube map texture, must have shape [minibatch_size, height, width, 3].
+        uv_da	(Optional) Tensor containing image-space derivatives of texture coordinates. Must have same shape as uv except for the last dimension that is to be twice as long.
+        mip_level_bias	(Optional) Per-pixel bias for mip level selection. If uv_da is omitted, determines mip level directly. Must have shape [minibatch_size, height, width].
+        mip	(Optional) Preconstructed mipmap stack from a texture_construct_mip() call, or a list of tensors specifying a custom mipmap stack. When specifying a custom mipmap stack, the tensors in the list must follow the same format as tex except for width and height that must follow the usual rules for mipmap sizes. The base level texture is still supplied in tex and must not be included in the list. Gradients of a custom mipmap stack are not automatically propagated to base texture but the mipmap tensors will receive gradients of their own. If a mipmap stack is not specified but the chosen filter mode requires it, the mipmap stack is constructed internally and discarded afterwards.
+        filter_mode	Texture filtering mode to be used. Valid values are 'auto', 'nearest', 'linear', 'linear-mipmap-nearest', and 'linear-mipmap-linear'. Mode 'auto' selects 'linear' if neither uv_da or mip_level_bias is specified, and 'linear-mipmap-linear' when at least one of them is specified, these being the highest-quality modes possible depending on the availability of the image-space derivatives of the texture coordinates or direct mip level information.
+        boundary_mode	Valid values are 'wrap', 'clamp', 'zero', and 'cube'. If tex defines a cube map, this must be set to 'cube'. The default mode 'wrap' takes fractional part of texture coordinates. Mode 'clamp' clamps texture coordinates to the centers of the boundary texels. Mode 'zero' virtually extends the texture with all-zero values in all directions.
+        max_mip_level	If specified, limits the number of mipmaps constructed and used in mipmap-based filter modes.
+        Returns:
+        A tensor containing the results of the texture sampling with shape [minibatch_size, height, width, tex_channels]. Cube map fetches with invalid uv coordinates (e.g., zero vectors) output all zeros and do not propagate gradients.
+        """
+        return dr.texture(
+            tex.float(),
+            uv.float(),
+            uv_da,
+            mip_level_bias,
+            mip,
+            filter_mode,
+            boundary_mode,
+            max_mip_level,
+        )
+
+    def antialias(
+        self, color, rast, pos, tri, topology_hash=None, pos_gradient_boost=1.0
+    ):
+        """
+        Perform antialiasing.
+
+        All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory.
+
+        Note that silhouette edge determination is based on vertex indices in the triangle tensor. For it to work properly, a vertex belonging to multiple triangles must be referred to using the same vertex index in each triangle. Otherwise, nvdiffrast will always classify the adjacent edges as silhouette edges, which leads to bad performance and potentially incorrect gradients. If you are unsure whether your data is good, check which pixels are modified by the antialias operation and compare to the example in the documentation.
+
+        Arguments:
+        color	Input image to antialias with shape [minibatch_size, height, width, num_channels].
+        rast	Main output tensor from rasterize().
+        pos	Vertex position tensor used in the rasterization operation.
+        tri	Triangle tensor used in the rasterization operation.
+        topology_hash	(Optional) Preconstructed topology hash for the triangle tensor. If not specified, the topology hash is constructed internally and discarded afterwards.
+        pos_gradient_boost	(Optional) Multiplier for gradients propagated to pos.
+        Returns:
+        A tensor containing the antialiased image with the same shape as color input tensor.
+        """
+        return dr.antialias(
+            color.float(),
+            rast,
+            pos.float(),
+            tri.int(),
+            topology_hash,
+            pos_gradient_boost,
+        )
+
+
+class DepthNormalizationStrategy(ABC):
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __call__(
+        self, depth: torch.FloatTensor, mask: torch.BoolTensor
+    ) -> torch.FloatTensor:
+        pass
+
+
+class DepthControlNetNormalization(DepthNormalizationStrategy):
+    def __init__(
+        self, far_clip: float = 0.25, near_clip: float = 1.0, bg_value: float = 0.0
+    ):
+        self.far_clip = far_clip
+        self.near_clip = near_clip
+        self.bg_value = bg_value
+
+    def __call__(
+        self, depth: torch.FloatTensor, mask: torch.BoolTensor
+    ) -> torch.FloatTensor:
+        batch_size = depth.shape[0]
+        min_depth = depth.view(batch_size, -1).min(dim=-1)[0][:, None, None]
+        max_depth = depth.view(batch_size, -1).max(dim=-1)[0][:, None, None]
+        depth = 1.0 - ((depth - min_depth) / (max_depth - min_depth + 1e-5)).clamp(
+            0.0, 1.0
+        )
+        depth = depth * (self.near_clip - self.far_clip) + self.far_clip
+        depth[~mask] = self.bg_value
+        return depth
+
+
+class Zero123PlusPlusNormalization(DepthNormalizationStrategy):
+    def __init__(self, bg_value: float = 0.8):
+        self.bg_value = bg_value
+
+    def __call__(self, depth: FloatTensor, mask: BoolTensor) -> FloatTensor:
+        batch_size = depth.shape[0]
+        min_depth = depth.view(batch_size, -1).min(dim=-1)[0][:, None, None]
+        max_depth = depth.view(batch_size, -1).max(dim=-1)[0][:, None, None]
+        depth = ((depth - min_depth) / (max_depth - min_depth + 1e-5)).clamp(0.0, 1.0)
+        depth[~mask] = self.bg_value
+        return depth
+
+
+class SimpleNormalization(DepthNormalizationStrategy):
+    def __init__(
+        self,
+        scale: float = 1.0,
+        offset: float = -1.0,
+        clamp: bool = True,
+        bg_value: float = 1.0,
+    ):
+        self.scale = scale
+        self.offset = offset
+        self.clamp = clamp
+        self.bg_value = bg_value
+
+    def __call__(self, depth: FloatTensor, mask: BoolTensor) -> FloatTensor:
+        depth = depth * self.scale + self.offset
+        if self.clamp:
+            depth = depth.clamp(0.0, 1.0)
+        depth[~mask] = self.bg_value
+        return depth
+
+
+def render(
+    ctx: NVDiffRastContextWrapper,
+    mesh: TexturedMesh,
+    cam: Camera,
+    height: int,
+    width: int,
+    render_attr: bool = True,
+    render_depth: bool = True,
+    render_normal: bool = True,
+    render_tangent: bool = False,
+    depth_normalization_strategy: DepthNormalizationStrategy = DepthControlNetNormalization(),
+    attr_background: Union[float, torch.FloatTensor] = 0.5,
+    antialias_attr=False,
+    normal_background: Union[float, torch.FloatTensor] = 0.0,
+    tangent_background: Union[float, torch.FloatTensor] = 0.0,
+    texture_override=None,
+    texture_filter_mode: str = "linear",
+) -> RenderOutput:
+    output_dict = {}
+
+    v_pos_clip = get_clip_space_position(mesh.v_pos, cam.mvp_mtx)
+    rast, _ = ctx.rasterize(v_pos_clip, mesh.t_pos_idx, (height, width), grad_db=True)
+    mask = rast[..., 3] > 0
+
+    gb_pos, _ = ctx.interpolate(mesh.v_pos[None], rast, mesh.t_pos_idx)
+    output_dict.update({"mask": mask, "pos": gb_pos})
+
+    if render_depth:
+        gb_pos_vs = transform_points_homo(gb_pos, cam.w2c)
+        gb_depth = -gb_pos_vs[..., 2]
+        # set background pixels to min depth value for correct min/max calculation
+        gb_depth = torch.where(
+            mask,
+            gb_depth,
+            gb_depth.view(gb_depth.shape[0], -1).min(dim=-1)[0][:, None, None],
+        )
+        gb_depth = depth_normalization_strategy(gb_depth, mask)
+        output_dict["depth"] = gb_depth
+
+    if render_attr:
+        tex_c, _ = ctx.interpolate(mesh.v_tex[None], rast, mesh.t_tex_idx)
+        texture = (
+            texture_override[None]
+            if texture_override is not None
+            else mesh.texture[None]
+        )
+        gb_rgb_fg = ctx.texture(texture, tex_c, filter_mode=texture_filter_mode)
+        gb_rgb_bg = torch.ones_like(gb_rgb_fg) * attr_background
+        gb_rgb = torch.where(mask[..., None], gb_rgb_fg, gb_rgb_bg)
+        if antialias_attr:
+            gb_rgb = ctx.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx)
+        output_dict["attr"] = gb_rgb
+
+    if render_normal:
+        gb_nrm, _ = ctx.interpolate(mesh.v_nrm[None], rast, mesh.stitched_t_pos_idx)
+        gb_nrm = F.normalize(gb_nrm, dim=-1, p=2)
+        gb_nrm[~mask] = normal_background
+        output_dict["normal"] = gb_nrm
+
+    if render_tangent:
+        gb_tang, _ = ctx.interpolate(mesh.v_tang[None], rast, mesh.stitched_t_pos_idx)
+        gb_tang = F.normalize(gb_tang, dim=-1, p=2)
+        gb_tang[~mask] = tangent_background
+        output_dict["tangent"] = gb_tang
+
+    return RenderOutput(**output_dict)
+
+
+def tensor_to_image(
+    data: Union[Image.Image, torch.Tensor, np.ndarray],
+    batched: bool = False,
+    format: str = "HWC",
+) -> Union[Image.Image, List[Image.Image]]:
+    if isinstance(data, Image.Image):
+        return data
+    if isinstance(data, torch.Tensor):
+        data = data.detach().cpu().numpy()
+    if data.dtype == np.float32 or data.dtype == np.float16:
+        data = (data * 255).astype(np.uint8)
+    elif data.dtype == np.bool_:
+        data = data.astype(np.uint8) * 255
+    assert data.dtype == np.uint8
+    if format == "CHW":
+        if batched and data.ndim == 4:
+            data = data.transpose((0, 2, 3, 1))
+        elif not batched and data.ndim == 3:
+            data = data.transpose((1, 2, 0))
+
+    if batched:
+        return [Image.fromarray(d) for d in data]
+    return Image.fromarray(data)

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/seg.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+
+import torch
+import transformers
+
+from .utils import IMAGE_TYPE, image_to_tensor
+
+
+class SegmentationModel(ABC):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __call__(self, images: IMAGE_TYPE) -> torch.FloatTensor:
+        pass
+
+
+class RMBGModel(SegmentationModel):
+    def __init__(self, pretrained_model_name_or_path: str, device: str):
+        self.model = transformers.AutoModelForImageSegmentation.from_pretrained(
+            pretrained_model_name_or_path, trust_remote_code=True
+        ).to(device)
+        self.device = device
+
+    def __call__(self, images: IMAGE_TYPE) -> torch.FloatTensor:
+        images = image_to_tensor(images, device=self.device)
+        batched = True
+        if images.ndim == 3:
+            images = images.unsqueeze(0)
+            batched = False
+
+        out = (
+            self.model(images.permute(0, 3, 1, 2) - 0.5)[0][0]
+            .clamp(0.0, 1.0)
+            .permute(0, 2, 3, 1)
+        )
+        if not batched:
+            out = out.squeeze(0)
+        return out

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/smart_paint.py

@@ -0,0 +1,334 @@
+import glob
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from itertools import product
+from typing import Any, Dict, List, Optional, Union
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from .camera import Camera, get_camera
+from .mesh import (
+    TexturedMesh,
+    load_mesh,
+    mesh_use_texture,
+    replace_mesh_texture_and_save,
+)
+from .projection import CameraProjection
+from .render import NVDiffRastContextWrapper, render
+from .utils import image_to_tensor, make_image_grid, tensor_to_image
+from .uv import (
+    ExponentialBlend,
+    NVDiffRastContextWrapper,
+    PoissonBlendingSolver,
+    SimpleUVValidityStrategy,
+    uv_blend,
+    uv_padding,
+    uv_precompute,
+    uv_render_attr,
+    uv_render_geometry,
+)
+
+
+class SmartPainter:
+    def __init__(self, device: str):
+        self.device = device
+        self.cam_proj = CameraProjection(
+            pb_backend="torch-cuda", bg_remover=None, device=device
+        )
+        self.ctx = NVDiffRastContextWrapper(device=self.device)
+
+    def __call__(
+        self,
+        mod_name: str,
+        mesh: TexturedMesh,
+        inpaint_func: callable,
+        uv_texture: torch.FloatTensor,
+        uv_inpaint_mask: torch.BoolTensor,
+        max_view_score_thresh: float = 0.02,
+        min_rounds: int = 3,
+        max_rounds: int = 8,
+        uv_padding_end: bool = True,
+        debug_dir: Optional[str] = None,
+        debug_visualize_details: bool = False,
+    ):
+        def make_view_selection_cam_params():
+            elevation_anchors = range(-60, 61, 15)
+            azimuth_anchors = range(0, 360, 30)
+            distance_anchors = [1.2]
+            fovy_deg_anchors = [40]
+            params = list(
+                product(
+                    elevation_anchors,
+                    azimuth_anchors,
+                    distance_anchors,
+                    fovy_deg_anchors,
+                )
+            )
+            params = list(zip(*params))
+            return params
+
+        elevation_deg, azimuth_deg, distance, fovy_deg = (
+            make_view_selection_cam_params()
+        )
+        view_selection_cams = get_camera(
+            elevation_deg=list(elevation_deg),
+            azimuth_deg=list(azimuth_deg),
+            distance=list(distance),
+            fovy_deg=list(fovy_deg),
+            perturb_camera_position=0.1,
+            device=self.device,
+        )
+
+        texture_update = uv_texture.clone()
+        uv_valid_mask_update = ~uv_inpaint_mask.clone()
+        score_map_update = torch.zeros_like(uv_valid_mask_update, dtype=torch.float32)
+        score_map_update[uv_valid_mask_update] = 1.0
+
+        max_view_score = 1
+        i = 0
+        while i < min_rounds or (
+            max_view_score > max_view_score_thresh and i < max_rounds
+        ):
+            score_map_image = score_map_update.float()[:, :, None].repeat_interleave(
+                3, dim=-1
+            )
+            render_size = 256
+            with mesh_use_texture(mesh, score_map_image):
+                render_out = render(
+                    self.ctx,
+                    mesh,
+                    view_selection_cams,
+                    height=render_size,
+                    width=render_size,
+                    attr_background=1.0,
+                    texture_filter_mode="nearest",
+                )
+
+            def get_view_aoi_cos(render_out, cameras):
+                view_normal = render_out.normal
+                view_normal_cs = (
+                    view_normal[:, :, :, None, :] * cameras.w2c[:, None, None, :3, :3]
+                ).sum(-1)
+                view_normal_cs = F.normalize(view_normal_cs, dim=-1, p=2)
+                view_normal_cs[~render_out.mask] = 0
+                view_aoi_cos = (
+                    (
+                        view_normal_cs
+                        * torch.as_tensor(
+                            [0.0, 0.0, 1.0],
+                            dtype=view_normal_cs.dtype,
+                            device=view_normal_cs.device,
+                        )[None, None, None]
+                    )
+                    .sum(-1)
+                    .clamp(0, 1)
+                )
+                return view_aoi_cos
+
+            view_aoi_cos = get_view_aoi_cos(render_out, view_selection_cams)
+            if debug_dir is not None and debug_visualize_details:
+                make_image_grid(tensor_to_image(render_out.attr, batched=True)).save(
+                    os.path.join(debug_dir, f"{mod_name}_render_score_{i:02d}.jpg")
+                )
+                make_image_grid(tensor_to_image(view_aoi_cos, batched=True)).save(
+                    os.path.join(
+                        debug_dir, f"{mod_name}_render_view_aoi_cos_{i:02d}.jpg"
+                    )
+                )
+            view_score = [
+                (
+                    ((attr < 1e-3) & (aoi_cos > 0.1)).sum().item()
+                    + (
+                        ((attr > 1e-3) & (aoi_cos > 0.1)).float()
+                        * (aoi_cos - attr - 0.3).clamp_min(0)
+                    )
+                    .sum()
+                    .item()
+                )
+                / render_size**2
+                for attr, aoi_cos in zip(render_out.attr[..., 0], view_aoi_cos)
+            ]
+
+            max_view_score = np.max(view_score)
+            best_view = np.argmax(view_score)
+
+            inpaint_render_size = 1024
+            best_cam = view_selection_cams[best_view : best_view + 1]
+
+            def shrink_mask(mask: torch.BoolTensor, radius: int) -> torch.BoolTensor:
+                return (
+                    (
+                        -torch.nn.functional.max_pool2d(
+                            -(mask[None, None].float()),
+                            kernel_size=radius * 2 + 1,
+                            stride=1,
+                            padding=radius,
+                        )
+                    )
+                    .squeeze()
+                    .bool()
+                )
+
+            def enlarge_mask(mask: torch.BoolTensor, radius: int) -> torch.BoolTensor:
+                return (
+                    torch.nn.functional.max_pool2d(
+                        mask[None, None].float(),
+                        kernel_size=radius * 2 + 1,
+                        stride=1,
+                        padding=radius,
+                    )
+                    .squeeze()
+                    .bool()
+                )
+
+            def blur_mask(mask: torch.BoolTensor, radius: int) -> torch.FloatTensor:
+                # Create a Gaussian kernel
+                x = torch.arange(
+                    -radius, radius + 1, dtype=torch.float32, device=mask.device
+                )
+                kernel = torch.exp(-(x**2) / (2 * (radius / 2) ** 2))
+                kernel = kernel[None, None, None]
+                kernel = kernel * kernel.transpose(2, 3)
+                kernel = kernel / kernel.sum()
+                # Apply convolution to blur the mask
+                blurred_mask = F.conv2d(
+                    mask[None, None].float(), kernel, padding=radius
+                )
+                return blurred_mask.squeeze()
+
+            def get_occulusion_boundary(
+                view_depth: torch.FloatTensor, dilation: int, thresh: float
+            ) -> torch.BoolTensor:
+                grad_x_kernel = torch.tensor(
+                    [[1, 0, -1], [2, 0, -2], [1, 0, -1]],
+                    device=view_depth.device,
+                    dtype=torch.float32,
+                ).view(1, 1, 3, 3)
+                grad_y_kernel = torch.tensor(
+                    [[1, 2, 1], [0, 0, 0], [-1, -2, -1]],
+                    device=view_depth.device,
+                    dtype=torch.float32,
+                ).view(1, 1, 3, 3)
+                view_depth_grad_x = F.conv2d(
+                    view_depth[None, None], grad_x_kernel, padding=1
+                )
+                view_depth_grad_y = F.conv2d(
+                    view_depth[None, None], grad_y_kernel, padding=1
+                )
+                view_depth_grad = (view_depth_grad_x**2 + view_depth_grad_y**2).sqrt()
+                occ_boundary = (view_depth_grad > thresh)[0][0]
+                if dilation > 0:
+                    occ_boundary = enlarge_mask(occ_boundary, dilation)
+                return occ_boundary
+
+            with mesh_use_texture(mesh, score_map_image):
+                render_out = render(
+                    self.ctx,
+                    mesh,
+                    best_cam,
+                    height=inpaint_render_size,
+                    width=inpaint_render_size,
+                    attr_background=1.0,
+                    texture_filter_mode="nearest",
+                )
+
+            view_aoi_cos = get_view_aoi_cos(render_out, best_cam)
+            inpaint_mask = (render_out.attr[0, :, :, 0] < 1e-3) | (
+                view_aoi_cos[0] - render_out.attr[0, :, :, 0] > 0.3
+            )
+            occ_boundary = get_occulusion_boundary(
+                render_out.depth[0], dilation=0, thresh=0.1
+            )
+            # shrink to remove uv seam
+            # enlarge the mask gives better results
+            # remove the occlusion boundary to avoid bleeding
+            inpaint_mask = enlarge_mask(shrink_mask(inpaint_mask, 3), 5) & ~occ_boundary
+            if debug_dir is not None and debug_visualize_details:
+                tensor_to_image(occ_boundary).save(
+                    os.path.join(debug_dir, f"{mod_name}_occ_boundary_{i:02d}.jpg")
+                )
+
+            with mesh_use_texture(mesh, texture_update):
+                inpaint_image = render(
+                    self.ctx,
+                    mesh,
+                    best_cam,
+                    height=inpaint_render_size,
+                    width=inpaint_render_size,
+                    texture_filter_mode="linear",
+                ).attr[0]
+
+            # according to https://github.com/advimman/lama/issues/262
+            # blur the sharp edge mask gives better results
+            # but it gives worse results in our case
+            # inpaint_mask = blur_mask(inpaint_mask, 1)
+
+            inpaint_mask_input = inpaint_mask.float()[None, None]
+            inpaint_image_input = inpaint_image.permute(2, 0, 1)[None]
+            inpaint_result = inpaint_func(inpaint_image_input, inpaint_mask_input)[
+                0
+            ].permute(1, 2, 0)
+
+            if debug_dir is not None:
+                make_image_grid(
+                    [
+                        tensor_to_image(inpaint_image),
+                        tensor_to_image(inpaint_mask),
+                        tensor_to_image(inpaint_result),
+                    ],
+                    rows=1,
+                ).save(
+                    os.path.join(debug_dir, f"{mod_name}_inpaint_result_{i:02d}.jpg")
+                )
+            with mesh_use_texture(mesh, texture_update):
+                proj_out = self.cam_proj(
+                    inpaint_result[None],
+                    mesh,
+                    best_cam,
+                    masks=inpaint_mask[None].float(),
+                    from_scratch=False,
+                    poisson_blending=False,
+                    depth_grad_dilation=3,
+                    uv_exp_blend_alpha=3,
+                    aoi_cos_valid_threshold=0.1,
+                    uv_size=mesh.uv_size,
+                    uv_padding=True,
+                    iou_rejection_threshold=None,
+                    return_dict=True,
+                )
+            uv_proj_inpaint, uv_valid_mask_inpaint = (
+                proj_out.uv_proj,
+                proj_out.uv_proj_mask,
+            )
+            if debug_dir is not None and debug_visualize_details:
+                tensor_to_image(uv_valid_mask_inpaint).save(
+                    os.path.join(debug_dir, f"{mod_name}_uv_inpaint_mask_{i:02d}.jpg")
+                )
+            texture_update = uv_proj_inpaint
+            uv_valid_mask_update = uv_valid_mask_inpaint | uv_valid_mask_update
+            score_map_inpaint = torch.where(
+                uv_valid_mask_inpaint,
+                proj_out.uv_aoi_cos[0],
+                torch.zeros_like(proj_out.uv_aoi_cos[0]),
+            )
+            score_map = torch.max(score_map_update, score_map_inpaint)
+            if debug_dir is not None and debug_visualize_details:
+                make_image_grid(
+                    [
+                        tensor_to_image(score_map_update),
+                        tensor_to_image(score_map_inpaint),
+                        tensor_to_image(score_map),
+                    ]
+                ).save(os.path.join(debug_dir, f"{mod_name}_score_map_{i:02d}.jpg"))
+            score_map_update = score_map
+            i += 1
+
+        if uv_padding_end:
+            texture_update = uv_padding(texture_update, uv_valid_mask_update, 3)
+
+        return texture_update, uv_valid_mask_update

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/utils.py

@@ -0,0 +1,139 @@
+import io
+import math
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional, Union
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+from PIL import Image
+from torch import BoolTensor, FloatTensor
+
+LIST_TYPE = Union[list, np.ndarray, torch.Tensor]
+IMAGE_TYPE = Union[Image.Image, List[Image.Image], np.ndarray, torch.Tensor]
+SINGLE_IMAGE_TYPE = Union[Image.Image, np.ndarray, torch.Tensor]
+
+
+def tensor_to_image(
+    data: Union[Image.Image, torch.Tensor, np.ndarray],
+    batched: bool = False,
+    format: str = "HWC",
+) -> Union[Image.Image, List[Image.Image]]:
+    if isinstance(data, Image.Image):
+        return data
+    if isinstance(data, torch.Tensor):
+        data = data.detach().cpu().numpy()
+    if data.dtype == np.float32 or data.dtype == np.float16:
+        data = (data * 255).astype(np.uint8)
+    elif data.dtype == np.bool_:
+        data = data.astype(np.uint8) * 255
+    assert data.dtype == np.uint8
+    if format == "CHW":
+        if batched and data.ndim == 4:
+            data = data.transpose((0, 2, 3, 1))
+        elif not batched and data.ndim == 3:
+            data = data.transpose((1, 2, 0))
+
+    if batched:
+        return [Image.fromarray(d) for d in data]
+    return Image.fromarray(data)
+
+
+def image_to_tensor(image: IMAGE_TYPE, return_type="pt", device: Optional[str] = None):
+    assert return_type in ["np", "pt"]
+    batched = True
+    if isinstance(image, Image.Image):
+        batched = False
+        image = [image]
+    if isinstance(image, list):
+        image = np.stack([np.array(img) for img in image], axis=0)
+        image = image.astype(np.float32) / 255.0
+    if isinstance(image, np.ndarray) and return_type == "pt":
+        image = torch.tensor(image, device=device)
+    if isinstance(image, torch.Tensor):
+        image = image.to(dtype=torch.float32, device=device)
+
+    if not batched:
+        image = image[0]
+    return image
+
+
+def largest_factor_near_sqrt(n: int) -> int:
+    """
+    Finds the largest factor of n that is closest to the square root of n.
+
+    Args:
+        n (int): The integer for which to find the largest factor near its square root.
+
+    Returns:
+        int: The largest factor of n that is closest to the square root of n.
+    """
+    sqrt_n = int(math.sqrt(n))  # Get the integer part of the square root
+
+    # First, check if the square root itself is a factor
+    if sqrt_n * sqrt_n == n:
+        return sqrt_n
+
+    # Otherwise, find the largest factor by iterating from sqrt_n downwards
+    for i in range(sqrt_n, 0, -1):
+        if n % i == 0:
+            return i
+
+    # If n is 1, return 1
+    return 1
+
+
+def make_image_grid(
+    images: List[Image.Image],
+    rows: Optional[int] = None,
+    cols: Optional[int] = None,
+    resize: Optional[int] = None,
+) -> Image.Image:
+    """
+    Prepares a single grid of images. Useful for visualization purposes.
+    """
+    if rows is None and cols is not None:
+        assert len(images) % cols == 0
+        rows = len(images) // cols
+    elif cols is None and rows is not None:
+        assert len(images) % rows == 0
+        cols = len(images) // rows
+    elif rows is None and cols is None:
+        rows = largest_factor_near_sqrt(len(images))
+        cols = len(images) // rows
+
+    assert len(images) == rows * cols
+
+    if resize is not None:
+        images = [img.resize((resize, resize)) for img in images]
+
+    w, h = images[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+
+    for i, img in enumerate(images):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+def get_current_timestamp(fmt: str = "%Y%m%d%H%M%S") -> str:
+    return datetime.now().strftime(fmt)
+
+
+def get_clip_space_position(pos: torch.FloatTensor, mvp_mtx: torch.FloatTensor):
+    pos_homo = torch.cat([pos, torch.ones([pos.shape[0], 1]).to(pos)], dim=-1)
+    return torch.matmul(pos_homo, mvp_mtx.permute(0, 2, 1))
+
+
+def transform_points_homo(pos: torch.FloatTensor, mtx: torch.FloatTensor):
+    batch_size = pos.shape[0]
+    pos_shape = pos.shape[1:-1]
+    pos = pos.reshape(batch_size, -1, 3)
+    pos_homo = torch.cat([pos, torch.ones_like(pos[..., 0:1])], dim=-1)
+    pos = (pos_homo.unsqueeze(2) * mtx.unsqueeze(1)).sum(-1)[..., :3]
+    pos = pos.reshape(batch_size, *pos_shape, 3)
+    return pos

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/uv.py

@@ -0,0 +1,468 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from .blend import PoissonBlendingSolver
+from .camera import Camera
+from .mesh import TexturedMesh
+from .render import NVDiffRastContextWrapper, SimpleNormalization, render
+from .utils import IMAGE_TYPE, get_clip_space_position, image_to_tensor
+
+
+@dataclass
+class UVPrecomputeOutput:
+    height: int
+    width: int
+    uv_attr: torch.Tensor
+    uv_mask: torch.Tensor
+    uv_pos: torch.Tensor
+
+
+def uv_precompute(
+    ctx: NVDiffRastContextWrapper,
+    mesh: TexturedMesh,
+    height: int,
+    width: int,
+    clone_attr: bool = False,
+) -> UVPrecomputeOutput:
+    uv_clip = mesh.v_tex * 2.0 - 1.0
+    uv_clip4 = torch.cat(
+        (
+            uv_clip,
+            torch.zeros_like(uv_clip[..., 0:1]),
+            torch.ones_like(uv_clip[..., 0:1]),
+        ),
+        dim=-1,
+    )
+    rast, _ = ctx.rasterize(uv_clip4[None], mesh.t_tex_idx, (height, width))
+    uv_mask = rast[0, :, :, 3] > 0
+
+    uv_pos, _ = ctx.interpolate(mesh.v_pos[None], rast, mesh.t_pos_idx)
+    uv_pos = uv_pos[0]
+
+    uv_precompute_output = UVPrecomputeOutput(
+        height=height,
+        width=width,
+        uv_attr=mesh.texture.clone() if clone_attr else mesh.texture,
+        uv_mask=uv_mask,
+        uv_pos=uv_pos,
+    )
+    return uv_precompute_output
+
+
+@dataclass
+class UVRenderGeometryOutput:
+    uv_pos_proj: torch.Tensor
+    uv_pos_error: torch.Tensor
+    uv_aoi_cos: torch.Tensor
+    uv_pos_ndc: torch.Tensor
+    view_mask: torch.Tensor
+    view_normal: torch.Tensor
+    view_aoi_cos: torch.Tensor
+    view_position: torch.Tensor
+    view_depth: torch.Tensor
+    view_depth_grad: Optional[torch.Tensor] = None
+    uv_depth_grad: Optional[torch.Tensor] = None
+    view_attr: Optional[torch.Tensor] = None
+
+
+def uv_render_geometry(
+    ctx: NVDiffRastContextWrapper,
+    mesh: TexturedMesh,
+    cam: Camera,
+    view_height: int,
+    view_width: int,
+    uv_precompute_output: UVPrecomputeOutput,
+    grid_sample_mode="bilinear",
+    compute_depth_grad: bool = False,
+    depth_grad_dilation: int = 1,
+    render_attr: bool = False,
+) -> UVRenderGeometryOutput:
+    batch_size = cam.c2w.shape[0]
+
+    height, width, _ = uv_precompute_output.uv_pos.shape
+    uv_pos_clip = get_clip_space_position(
+        uv_precompute_output.uv_pos.view(-1, 3), cam.mvp_mtx
+    ).view(batch_size, height, width, 4)
+    uv_pos_ndc = uv_pos_clip[..., :2] / uv_pos_clip[..., 3:4]
+
+    render_output = render(
+        ctx,
+        mesh,
+        cam,
+        view_height,
+        view_width,
+        render_attr=render_attr,
+        render_depth=True,
+        render_normal=True,
+        depth_normalization_strategy=SimpleNormalization(
+            scale=1.0, offset=0.0, clamp=False, bg_value=1e2
+        ),
+    )
+    view_position = render_output.pos
+    view_mask = render_output.mask
+    view_normal = render_output.normal
+    view_normal_cs = (
+        view_normal[:, :, :, None, :] * cam.w2c[:, None, None, :3, :3]
+    ).sum(-1)
+    view_normal_cs = F.normalize(view_normal_cs, dim=-1, p=2)
+    view_normal_cs[~render_output.mask] = render_output.normal[~render_output.mask]
+    view_aoi_cos = (
+        view_normal_cs
+        * torch.as_tensor(
+            [0.0, 0.0, 1.0], dtype=view_normal_cs.dtype, device=view_normal_cs.device
+        )[None, None, None]
+    ).sum(-1)
+    view_aoi_cos = view_aoi_cos.clamp(0.0, 1.0)
+
+    view_depth = render_output.depth
+    if compute_depth_grad:
+        grad_x_kernel = torch.tensor(
+            [[1, 0, -1], [2, 0, -2], [1, 0, -1]],
+            device=view_depth.device,
+            dtype=torch.float32,
+        ).view(1, 1, 3, 3)
+        grad_y_kernel = torch.tensor(
+            [[1, 2, 1], [0, 0, 0], [-1, -2, -1]],
+            device=view_depth.device,
+            dtype=torch.float32,
+        ).view(1, 1, 3, 3)
+        view_depth_grad_x = F.conv2d(view_depth[:, None], grad_x_kernel, padding=1)
+        view_depth_grad_y = F.conv2d(view_depth[:, None], grad_y_kernel, padding=1)
+        view_depth_grad = (view_depth_grad_x**2 + view_depth_grad_y**2).sqrt()
+        view_depth_grad = F.max_pool2d(
+            view_depth_grad,
+            depth_grad_dilation,
+            stride=1,
+            padding=depth_grad_dilation // 2,
+        )
+
+        uv_depth_grad = F.grid_sample(
+            view_depth_grad, uv_pos_ndc, align_corners=False, mode=grid_sample_mode
+        ).permute(0, 2, 3, 1)[..., 0]
+    else:
+        view_depth_grad = None
+        uv_depth_grad = None
+
+    if render_attr:
+        view_attr = render_output.attr
+    else:
+        view_attr = None
+
+    uv_pos_proj = F.grid_sample(
+        view_position.permute(0, 3, 1, 2),
+        uv_pos_ndc,
+        align_corners=False,
+        mode=grid_sample_mode,
+    ).permute(0, 2, 3, 1)
+
+    uv_pos_error = ((uv_pos_proj - uv_precompute_output.uv_pos) ** 2).sum(-1).sqrt()
+
+    uv_aoi_cos = F.grid_sample(
+        view_aoi_cos[..., None].permute(0, 3, 1, 2),
+        uv_pos_ndc,
+        align_corners=False,
+        mode=grid_sample_mode,
+    ).permute(0, 2, 3, 1)[..., 0]
+
+    return UVRenderGeometryOutput(
+        uv_pos_proj=uv_pos_proj,
+        uv_pos_error=uv_pos_error,
+        uv_aoi_cos=uv_aoi_cos,
+        uv_pos_ndc=uv_pos_ndc,
+        view_mask=view_mask,
+        view_position=view_position,
+        view_normal=view_normal,
+        view_aoi_cos=view_aoi_cos,
+        view_depth=view_depth,
+        view_depth_grad=view_depth_grad,
+        uv_depth_grad=uv_depth_grad,
+        view_attr=view_attr,
+    )
+
+
+@dataclass
+class UVRenderAttrOutput:
+    uv_attr_proj: torch.Tensor
+    uv_mask_proj: Optional[torch.Tensor]
+
+
+def uv_render_attr(
+    images: IMAGE_TYPE,
+    uv_render_geometry_output: UVRenderGeometryOutput,
+    masks: Optional[IMAGE_TYPE] = None,
+    grid_sample_mode: str = "bilinear",
+) -> UVRenderAttrOutput:
+    images = image_to_tensor(images, device=uv_render_geometry_output.uv_pos_ndc.device)
+    uv_attr_proj = F.grid_sample(
+        images.permute(0, 3, 1, 2),
+        uv_render_geometry_output.uv_pos_ndc,
+        align_corners=False,
+        mode=grid_sample_mode,
+    ).permute(0, 2, 3, 1)
+
+    if masks is not None:
+        masks = image_to_tensor(
+            masks, device=uv_render_geometry_output.uv_pos_ndc.device
+        )
+        if masks.ndim == 4:
+            masks = masks.mean(-1)
+        uv_mask_proj = F.grid_sample(
+            masks[..., None].permute(0, 3, 1, 2),
+            uv_render_geometry_output.uv_pos_ndc,
+            align_corners=False,
+            mode=grid_sample_mode,
+        ).permute(0, 2, 3, 1)[..., 0]
+    else:
+        uv_mask_proj = None
+
+    return UVRenderAttrOutput(uv_attr_proj=uv_attr_proj, uv_mask_proj=uv_mask_proj)
+
+
+@dataclass
+class UVBlendOutput:
+    uv_attr_blend: Optional[torch.Tensor]
+    uv_valid_mask: torch.Tensor
+    uv_valid_mask_blend: torch.Tensor
+    uv_blend_weight: torch.Tensor
+
+
+class UVValidityStrategy(ABC):
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __call__(
+        self,
+        uv_precompute_output: UVPrecomputeOutput,
+        uv_render_geometry_output: UVRenderGeometryOutput,
+        uv_render_attr_output: Optional[UVRenderAttrOutput],
+    ) -> torch.BoolTensor:
+        pass
+
+
+class SimpleUVValidityStrategy(UVValidityStrategy):
+    def __init__(
+        self,
+        pos_error_eps: float = 1e-3,
+        aoi_cos_thresh: float = 0.1,
+        mask_thresh: float = 0.9,
+        depth_grad_thresh: Optional[float] = None,
+        first_view_dominate: bool = False,
+    ):
+        self.pos_error_eps = pos_error_eps
+        self.aoi_cos_thresh = aoi_cos_thresh
+        self.mask_thresh = mask_thresh
+        self.depth_grad_thresh = depth_grad_thresh
+        self.first_view_dominate = first_view_dominate
+
+    def __call__(
+        self,
+        uv_precompute_output: UVPrecomputeOutput,
+        uv_render_geometry_output: UVRenderGeometryOutput,
+        uv_render_attr_output: Optional[UVRenderAttrOutput],
+    ) -> torch.BoolTensor:
+        uv_pos_valid = uv_render_geometry_output.uv_pos_error < self.pos_error_eps
+        uv_aoi_cos_valid = uv_render_geometry_output.uv_aoi_cos > self.aoi_cos_thresh
+        uv_valid_mask = uv_pos_valid & uv_aoi_cos_valid
+
+        if self.depth_grad_thresh is not None:
+            if uv_render_geometry_output.uv_depth_grad is None:
+                print(
+                    "Warning: Depth gradient is not computed, depth gradient threshold is ignored."
+                )
+            else:
+                uv_valid_mask &= (
+                    uv_render_geometry_output.uv_depth_grad < self.depth_grad_thresh
+                )
+
+        uv_valid_mask &= uv_precompute_output.uv_mask
+
+        if (
+            uv_render_attr_output is not None
+            and uv_render_attr_output.uv_mask_proj is not None
+        ):
+            uv_valid_mask &= uv_render_attr_output.uv_mask_proj > self.mask_thresh
+        else:
+            print("No view mask provided for UV blending, using all valid pixels")
+
+        if self.first_view_dominate:
+            uv_valid_mask[1:][
+                uv_valid_mask[0:1].expand(uv_valid_mask.shape[0] - 1, -1, -1)
+            ] = False
+
+        return uv_valid_mask
+
+
+class UVBlendWeightStrategy(ABC):
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def __call__(
+        self,
+        uv_precompute_output: UVPrecomputeOutput,
+        uv_render_geometry_output: UVRenderGeometryOutput,
+        uv_render_attr_output: Optional[UVRenderAttrOutput],
+        uv_valid_mask: torch.BoolTensor,
+    ) -> torch.FloatTensor:
+        pass
+
+
+class ExponentialBlend(UVBlendWeightStrategy):
+    def __init__(
+        self,
+        alpha: float = 1.0,
+        normalization: str = "linear",
+        view_weight: Optional[torch.FloatTensor] = None,
+    ):
+        self.alpha = alpha
+        self.normalization = normalization
+        self.view_weight = view_weight
+
+    def __call__(
+        self,
+        uv_precompute_output: UVPrecomputeOutput,
+        uv_render_geometry_output: UVRenderGeometryOutput,
+        uv_render_attr_output: Optional[UVRenderAttrOutput],
+        uv_valid_mask: torch.BoolTensor,
+    ) -> torch.FloatTensor:
+        weight = uv_render_geometry_output.uv_aoi_cos
+        weight = weight * uv_valid_mask.float()
+        if self.view_weight is not None:
+            weight = weight ** (self.alpha / self.view_weight[:, None, None].to(weight))
+        else:
+            weight = weight**self.alpha
+        if self.normalization == "linear":
+            weight_norm = (weight / weight.sum(axis=0, keepdim=True).clamp(1e-5)).clamp(
+                0.0, 1.0
+            )
+        elif self.normalization == "softmax":
+            weight[~uv_valid_mask] = -1e5
+            weight_norm = F.softmax(weight, dim=0)
+        return weight_norm
+
+
+class RandomChoiceBlend(UVBlendWeightStrategy):
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def __call__(
+        self,
+        uv_precompute_output: UVPrecomputeOutput,
+        uv_render_geometry_output: UVRenderGeometryOutput,
+        uv_render_attr_output: Optional[UVRenderAttrOutput],
+        uv_valid_mask: torch.BoolTensor,
+    ) -> torch.FloatTensor:
+        weight = uv_render_geometry_output.uv_aoi_cos
+        weight = weight * uv_valid_mask.float()
+        weight[weight > 0] = torch.rand_like(weight[weight > 0])
+        weight_norm = (
+            F.one_hot(weight.max(dim=0).indices, num_classes=weight.shape[0])
+            .float()
+            .permute(2, 0, 1)
+        )
+        return weight_norm
+
+
+def uv_padding(attr: torch.FloatTensor, inside_mask: torch.BoolTensor, radius: int):
+    if attr.min() < -1e-5 or attr.max() > 1 + 1e-5:
+        print(
+            "Warning: UV blending result is out of range [0, 1], force clamped before UV padding."
+        )
+
+    from .cv_ops import inpaint_torch
+
+    attr_padded = inpaint_torch(attr.clamp(0.0, 1.0), ~inside_mask, radius)
+    return attr_padded
+
+
+def uv_blend(
+    uv_precompute_output: UVPrecomputeOutput,
+    uv_render_geometry_output: UVRenderGeometryOutput,
+    uv_render_attr_output: Optional[UVRenderAttrOutput],
+    uv_validity_strategy: UVValidityStrategy = SimpleUVValidityStrategy(),
+    uv_blend_weight_strategy: UVBlendWeightStrategy = ExponentialBlend(),
+    empty_value: float = 0.0,
+    do_uv_padding: bool = True,
+    uv_padding_radius: int = 3,
+    pad_unseen_area: bool = False,
+    poisson_blending: bool = False,
+    pb_solver: Optional[PoissonBlendingSolver] = None,
+    pb_num_iters: int = 1000,
+    pb_keep_original_border: bool = True,
+    pb_inplace: bool = False,
+    pb_grad_mode: str = "src",
+) -> UVBlendOutput:
+    uv_valid_mask = uv_validity_strategy(
+        uv_precompute_output, uv_render_geometry_output, uv_render_attr_output
+    )
+    uv_blend_weight = uv_blend_weight_strategy(
+        uv_precompute_output,
+        uv_render_geometry_output,
+        uv_render_attr_output,
+        uv_valid_mask,
+    )
+    uv_valid_mask_blend = uv_valid_mask.any(dim=0)
+
+    if uv_render_attr_output is None:
+        return UVBlendOutput(
+            uv_attr_blend=None,
+            uv_valid_mask=uv_valid_mask,
+            uv_valid_mask_blend=uv_valid_mask_blend,
+            uv_blend_weight=uv_blend_weight,
+        )
+
+    uv_attr_blend = (
+        uv_render_attr_output.uv_attr_proj * uv_blend_weight[..., None]
+    ).sum(axis=0)
+    if poisson_blending:
+        assert do_uv_padding
+        assert pb_solver is not None
+        uv_attr_blend_padded = uv_padding(
+            uv_attr_blend, uv_valid_mask_blend, uv_padding_radius
+        )
+        if pb_keep_original_border:
+            pb_tgt = uv_precompute_output.uv_attr
+        else:
+            uv_attr_hard_stitch = (
+                uv_attr_blend * uv_valid_mask_blend[..., None].float()
+                + uv_precompute_output.uv_attr
+                * (~uv_valid_mask_blend)[..., None].float()
+            )
+            uv_attr_hard_stitch_padded = uv_padding(
+                uv_attr_hard_stitch, uv_precompute_output.uv_mask, uv_padding_radius
+            )
+            pb_tgt = uv_attr_hard_stitch_padded
+
+        uv_attr_blend = pb_solver(
+            uv_attr_blend_padded,
+            uv_valid_mask_blend,
+            pb_tgt,
+            pb_num_iters,
+            inplace=pb_inplace,
+            grad_mode=pb_grad_mode,
+        )
+    else:
+        uv_attr_blend = (
+            uv_attr_blend * uv_valid_mask_blend[..., None].float()
+            + uv_precompute_output.uv_attr * (~uv_valid_mask_blend)[..., None].float()
+        )
+
+    if do_uv_padding:
+        content_mask = (
+            uv_valid_mask_blend if pad_unseen_area else uv_precompute_output.uv_mask
+        )
+        uv_attr_blend = uv_padding(uv_attr_blend, content_mask, uv_padding_radius)
+
+    return UVBlendOutput(
+        uv_attr_blend=uv_attr_blend,
+        uv_valid_mask=uv_valid_mask,
+        uv_valid_mask_blend=uv_valid_mask_blend,
+        uv_blend_weight=uv_blend_weight,
+    )

Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/warp.py

@@ -0,0 +1,285 @@
+import os
+
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+from einops import rearrange
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+
+
+def draw(img, lines):
+    """
+    lines: n x 4, x1,y2,x2,y2
+    """
+    if isinstance(img, torch.Tensor):
+        img = img.detach().cpu().numpy()
+
+    if isinstance(img, np.ndarray):
+        if img.dtype == np.float32 or img.dtype == np.float64:
+            img = (img * 255.0).astype(np.uint8)
+        img = Image.fromarray(img)
+
+    img = img.resize((512, 512), Image.Resampling.BICUBIC)
+    w, h = img.size
+    assert w == h
+    if isinstance(lines, torch.Tensor):
+        lines = lines.detach().cpu().numpy()
+    lines = (lines + 1) / 2 * w
+
+    img1 = ImageDraw.Draw(img)
+    for line in lines:
+        img1.line(line.tolist(), fill="red", width=1)
+    return img
+
+
+def construct_grid_mesh(n_grid):
+    vertices = []
+    vertices_movable_ids = []
+    idx = 0
+    for j in range(n_grid + 1):
+        for i in range(n_grid + 1):
+            # 3D fit in trimesh format
+            if 0 < i < n_grid and 0 < j < n_grid:
+                vertices_movable_ids.append(idx)
+            vertices.append([i / n_grid, j / n_grid, 1.0 / 2])
+            idx += 1
+
+    vertices = np.array(vertices)
+    vertices = 2 * vertices - 1
+
+    faces = []
+    for j in range(n_grid):
+        for i in range(n_grid):
+            # clockwise
+            faces.append(
+                [
+                    i + j * (n_grid + 1),
+                    i + 1 + j * (n_grid + 1),
+                    i + (j + 1) * (n_grid + 1),
+                ]
+            )
+            faces.append(
+                [
+                    i + 1 + j * (n_grid + 1),
+                    i + 1 + (j + 1) * (n_grid + 1),
+                    i + (j + 1) * (n_grid + 1),
+                ]
+            )
+    faces = np.array(faces)
+    vertices_movable_ids = np.array(vertices_movable_ids)
+
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+    return mesh, vertices_movable_ids, vertices
+
+
+def compute_warp_field(
+    ctx,
+    src_images_tensor,
+    tgt_images_tensor,
+    n_grid,
+    optim_res,
+    optim_step_per_res,
+    lambda_reg,
+    temp_dir,
+    verbose,
+    device,
+):
+    """
+    src_images_tensor: 4 H W 4
+    tgt_images_tensor: 4 H W 4
+    return: 4 H W 4
+    """
+    lam_reg = lambda_reg
+    lam_mask = 2
+    mesh, _, vertices_np = construct_grid_mesh(n_grid)
+
+    # vertices = mesh.vertices
+    faces = mesh.faces
+    edges = mesh.edges_unique
+
+    n_degree = torch.tensor(mesh.vertex_degree)
+    vertices = torch.tensor(vertices_np, device=device, dtype=torch.float32)
+    vertices_unopt = vertices.clone()
+    faces = torch.tensor(faces, device=device)
+    edges = torch.tensor(edges, device=device)
+
+    warped_images_tensor = []
+    bs = src_images_tensor.shape[0]
+
+    for img_idx in range(bs):
+        src_image_tensor = src_images_tensor[img_idx]
+        tgt_image_tensor = tgt_images_tensor[img_idx]
+
+        if verbose and temp_dir is not None:
+            vis_dir = os.path.join(temp_dir, f"{img_idx}")
+            os.makedirs(vis_dir, exist_ok=True)
+        # prepare for nvdiff rendering
+        v_origin_homo = torch.cat(
+            [vertices_unopt, torch.ones_like(vertices_unopt[..., :1])], dim=-1
+        )
+
+        # prepare parameters
+        v_move_indices = torch.where(n_degree == 6)[0]
+        vertices_movable = vertices[v_move_indices].detach().clone()
+        vertices_movable.requires_grad = True
+        opt = torch.optim.Adam([vertices_movable], lr=0.02)
+
+        for resl in optim_res:
+            rast, _ = dr.rasterize(
+                ctx,
+                v_origin_homo[None, ...].float(),
+                faces.int(),
+                (resl, resl),
+                grad_db=True,
+            )
+            face_ids = rast[..., 3].long() - 1
+            assert torch.all(face_ids >= 0)
+            u = rast[..., 0]
+            v = rast[..., 1]
+            pixel_vertice_ids = faces[face_ids]
+            pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
+
+            src_img = F.interpolate(
+                rearrange(src_image_tensor[..., :3], "H W C -> () C H W"),
+                size=(resl, resl),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            tgt_img = F.interpolate(
+                rearrange(tgt_image_tensor[..., :3], "H W C -> () C H W"),
+                size=(resl, resl),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # src_img, tgt_img: 1 C H W
+            if verbose:
+                Image.fromarray(
+                    (tgt_img[0].permute(1, 2, 0).detach().cpu().numpy() * 255.0).astype(
+                        np.uint8
+                    )
+                ).save(os.path.join(vis_dir, f"target_{resl:04d}.png"))
+            with tqdm(range(optim_step_per_res), disable=not verbose) as pbar:
+                for i in pbar:
+                    opt.zero_grad()
+                    vertices_all = vertices_unopt.detach().clone()
+                    vertices_all[v_move_indices] = vertices_movable
+
+                    pixel_vertices = vertices_all[
+                        pixel_vertice_ids
+                    ]  # 1 x 512 x 512 x 3 x 3
+                    pixel_coords_inter = torch.sum(
+                        pixel_vertices * pixel_bary_coords[..., None], dim=-2
+                    )
+                    src_img_warped = F.grid_sample(
+                        src_img,
+                        pixel_coords_inter[..., :2],
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+
+                    edge_vertices_all = vertices_all[edges]
+                    edge_vertices_unopt = vertices_unopt[edges]
+
+                    edge_len_all = torch.linalg.norm(
+                        edge_vertices_all[:, 0, :2] - edge_vertices_all[:, 1, :2],
+                        dim=-1,
+                    )
+                    edge_len_unopt = torch.linalg.norm(
+                        edge_vertices_unopt[:, 0, :2] - edge_vertices_all[:, 1, :2],
+                        dim=-1,
+                    )
+
+                    reg_loss = ((edge_len_all - edge_len_unopt) ** 2).mean()
+
+                    # mask_loss = ((src_img_warped[:, 3:4, :, :] - tgt_img[:, 3:4, :, :])**2).mean()
+                    # rgb_loss = (((src_img_warped[:, :3, :, :] - tgt_img[:, :3, :, :]) * tgt_img[:, 3:4, :, :])**2).mean()
+                    # loss = rgb_loss + lam_reg * reg_loss + lam_mask * mask_loss
+                    img_loss = ((src_img_warped - tgt_img) ** 2).mean()
+                    loss = img_loss + lam_reg * reg_loss
+
+                    loss.backward()
+                    opt.step()
+                    if verbose:
+                        print(img_loss.item(), reg_loss.item())
+
+                    if verbose:
+                        draw(
+                            src_img[0].permute(1, 2, 0),
+                            edge_vertices_all[:, :, :2].reshape(-1, 4),
+                        ).save(os.path.join(vis_dir, f"src_{resl:04d}_{i:03d}.png"))
+                        Image.fromarray(
+                            (
+                                torch.cat(
+                                    [
+                                        tgt_img,
+                                        src_img_warped,
+                                        (tgt_img - src_img_warped).abs(),
+                                    ],
+                                    dim=-1,
+                                )[0]
+                                .permute(1, 2, 0)
+                                .detach()
+                                .cpu()
+                                .numpy()
+                                * 255.0
+                            ).astype(np.uint8)
+                        ).save(os.path.join(vis_dir, f"opt_{resl:04d}_{i:03d}.png"))
+                        warped_this_resl = Image.fromarray(
+                            (
+                                src_img_warped[0]
+                                .permute(1, 2, 0)
+                                .detach()
+                                .cpu()
+                                .numpy()
+                                * 255.0
+                            ).astype(np.uint8)
+                        )
+                        warped_this_resl.save(
+                            os.path.join(vis_dir, f"warped_{resl:04d}_{i:03d}.png")
+                        )
+
+        # warp on 512 resolution
+        resl = src_image_tensor.shape[1]
+        with torch.no_grad():
+            rast, _ = dr.rasterize(
+                ctx,
+                v_origin_homo[None, ...].float(),
+                faces.int(),
+                (resl, resl),
+                grad_db=True,
+            )
+            face_ids = rast[..., 3].long() - 1
+            assert torch.all(face_ids >= 0)
+            u = rast[..., 0]
+            v = rast[..., 1]
+            pixel_vertice_ids = faces[face_ids]
+            pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
+
+            vertices_all = vertices_unopt.detach().clone()
+            vertices_all[v_move_indices] = vertices_movable
+
+            pixel_vertices = vertices_all[pixel_vertice_ids]  # 1 x 512 x 512 x 3 x 3
+            pixel_coords_inter = torch.sum(
+                pixel_vertices * pixel_bary_coords[..., None], dim=-2
+            )
+            src_img_warped = rearrange(
+                F.grid_sample(
+                    rearrange(src_image_tensor, "H W C -> () C H W"),
+                    pixel_coords_inter[..., :2],
+                    mode="bicubic",
+                    align_corners=False,
+                ),
+                "() C H W -> H W C",
+            ).clamp(0, 1)
+        warped_images_tensor.append(src_img_warped)
+
+    warped_images_tensor = torch.stack(warped_images_tensor, dim=0)  # 4 H W 3
+
+    return warped_images_tensor

Gen_3D_Modules/MV_Adapter/mvadapter/utils/misc.py

@@ -0,0 +1,220 @@
+import gc
+import os
+import re
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+
+import psutil
+import torch
+from packaging import version
+
+from .config import config_to_primitive
+from .core import debug, find, info, warn
+from .typing import *
+
+
+def parse_version(ver: str):
+    return version.parse(ver)
+
+
+def get_rank():
+    # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
+    # therefore LOCAL_RANK needs to be checked first
+    rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
+    for key in rank_keys:
+        rank = os.environ.get(key)
+        if rank is not None:
+            return int(rank)
+    return 0
+
+
+def get_device():
+    return torch.device(f"cuda:{get_rank()}")
+
+
+def load_module_weights(
+    path, module_name=None, ignore_modules=None, mapping=None, map_location=None
+) -> Tuple[dict, int, int]:
+    if module_name is not None and ignore_modules is not None:
+        raise ValueError("module_name and ignore_modules cannot be both set")
+    if map_location is None:
+        map_location = get_device()
+
+    ckpt = torch.load(path, map_location=map_location)
+    state_dict = ckpt["state_dict"]
+
+    if mapping is not None:
+        state_dict_to_load = {}
+        for k, v in state_dict.items():
+            if any([k.startswith(m["to"]) for m in mapping]):
+                pass
+            else:
+                state_dict_to_load[k] = v
+        for k, v in state_dict.items():
+            for m in mapping:
+                if k.startswith(m["from"]):
+                    k_dest = k.replace(m["from"], m["to"])
+                    info(f"Mapping {k} => {k_dest}")
+                    state_dict_to_load[k_dest] = v.clone()
+        state_dict = state_dict_to_load
+
+    state_dict_to_load = state_dict
+
+    if ignore_modules is not None:
+        state_dict_to_load = {}
+        for k, v in state_dict.items():
+            ignore = any(
+                [k.startswith(ignore_module + ".") for ignore_module in ignore_modules]
+            )
+            if ignore:
+                continue
+            state_dict_to_load[k] = v
+
+    if module_name is not None:
+        state_dict_to_load = {}
+        for k, v in state_dict.items():
+            m = re.match(rf"^{module_name}\.(.*)$", k)
+            if m is None:
+                continue
+            state_dict_to_load[m.group(1)] = v
+
+    return state_dict_to_load, ckpt["epoch"], ckpt["global_step"]
+
+
+def C(value: Any, epoch: int, global_step: int) -> float:
+    if isinstance(value, int) or isinstance(value, float):
+        pass
+    else:
+        value = config_to_primitive(value)
+        if not isinstance(value, list):
+            raise TypeError("Scalar specification only supports list, got", type(value))
+        if len(value) == 3:
+            value = [0] + value
+        assert len(value) == 4
+        start_step, start_value, end_value, end_step = value
+        if isinstance(end_step, int):
+            current_step = global_step
+            value = start_value + (end_value - start_value) * max(
+                min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
+            )
+        elif isinstance(end_step, float):
+            current_step = epoch
+            value = start_value + (end_value - start_value) * max(
+                min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
+            )
+    return value
+
+
+def cleanup():
+    gc.collect()
+    torch.cuda.empty_cache()
+    try:
+        import tinycudann as tcnn
+
+        tcnn.free_temporary_memory()
+    except:
+        pass
+
+
+def finish_with_cleanup(func: Callable):
+    def wrapper(*args, **kwargs):
+        out = func(*args, **kwargs)
+        cleanup()
+        return out
+
+    return wrapper
+
+
+def _distributed_available():
+    return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+
+def barrier():
+    if not _distributed_available():
+        return
+    else:
+        torch.distributed.barrier()
+
+
+def broadcast(tensor, src=0):
+    if not _distributed_available():
+        return tensor
+    else:
+        torch.distributed.broadcast(tensor, src=src)
+        return tensor
+
+
+def enable_gradient(model, enabled: bool = True) -> None:
+    for param in model.parameters():
+        param.requires_grad_(enabled)
+
+
+class TimeRecorder:
+    _instance = None
+
+    def __init__(self):
+        self.items = {}
+        self.accumulations = defaultdict(list)
+        self.time_scale = 1000.0  # ms
+        self.time_unit = "ms"
+        self.enabled = False
+
+    def __new__(cls):
+        # singleton
+        if cls._instance is None:
+            cls._instance = super(TimeRecorder, cls).__new__(cls)
+        return cls._instance
+
+    def enable(self, enabled: bool) -> None:
+        self.enabled = enabled
+
+    def start(self, name: str) -> None:
+        if not self.enabled:
+            return
+        torch.cuda.synchronize()
+        self.items[name] = time.time()
+
+    def end(self, name: str, accumulate: bool = False) -> float:
+        if not self.enabled or name not in self.items:
+            return
+        torch.cuda.synchronize()
+        start_time = self.items.pop(name)
+        delta = time.time() - start_time
+        if accumulate:
+            self.accumulations[name].append(delta)
+        t = delta * self.time_scale
+        info(f"{name}: {t:.2f}{self.time_unit}")
+
+    def get_accumulation(self, name: str, average: bool = False) -> float:
+        if not self.enabled or name not in self.accumulations:
+            return
+        acc = self.accumulations.pop(name)
+        total = sum(acc)
+        if average:
+            t = total / len(acc) * self.time_scale
+        else:
+            t = total * self.time_scale
+        info(f"{name} for {len(acc)} times: {t:.2f}{self.time_unit}")
+
+
+### global time recorder
+time_recorder = TimeRecorder()
+
+
+@contextmanager
+def time_recorder_enabled():
+    enabled = time_recorder.enabled
+    time_recorder.enable(enabled=True)
+    try:
+        yield
+    finally:
+        time_recorder.enable(enabled=enabled)
+
+
+def show_vram_usage(name):
+    available, total = torch.cuda.mem_get_info()
+    used = total - available
+    print(
+        f"{name}: {used / 1024**2:.1f}MB, {psutil.Process(os.getpid()).memory_info().rss / 1024**2:.1f}MB"
+    )

Gen_3D_Modules/MV_Adapter/mvadapter/utils/ops.py

@@ -0,0 +1,462 @@
+import math
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+from .core import debug, find, info, warn
+from .typing import *
+
+
+def dot(x, y):
+    return torch.sum(x * y, -1, keepdim=True)
+
+
+def reflect(x, n):
+    return 2 * dot(x, n) * n - x
+
+
+ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]]
+
+
+def scale_tensor(
+    dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale
+):
+    if inp_scale is None:
+        inp_scale = (0, 1)
+    if tgt_scale is None:
+        tgt_scale = (0, 1)
+    if isinstance(tgt_scale, Tensor):
+        assert dat.shape[-1] == tgt_scale.shape[-1]
+    dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
+    dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
+    return dat
+
+
+class _TruncExp(Function):  # pylint: disable=abstract-method
+    # Implementation from torch-ngp:
+    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, x):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(x)
+        return torch.exp(x)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, g):  # pylint: disable=arguments-differ
+        x = ctx.saved_tensors[0]
+        return g * torch.exp(torch.clamp(x, max=15))
+
+
+trunc_exp = _TruncExp.apply
+
+
+def get_activation(name) -> Callable:
+    if name is None:
+        return lambda x: x
+    name = name.lower()
+    if name == "none":
+        return lambda x: x
+    elif name == "lin2srgb":
+        return lambda x: torch.where(
+            x > 0.0031308,
+            torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
+            12.92 * x,
+        ).clamp(0.0, 1.0)
+    elif name == "exp":
+        return lambda x: torch.exp(x)
+    elif name == "shifted_exp":
+        return lambda x: torch.exp(x - 1.0)
+    elif name == "trunc_exp":
+        return trunc_exp
+    elif name == "shifted_trunc_exp":
+        return lambda x: trunc_exp(x - 1.0)
+    elif name == "sigmoid":
+        return lambda x: torch.sigmoid(x)
+    elif name == "tanh":
+        return lambda x: torch.tanh(x)
+    elif name == "shifted_softplus":
+        return lambda x: F.softplus(x - 1.0)
+    elif name == "scale_-11_01":
+        return lambda x: x * 0.5 + 0.5
+    elif name == "negative":
+        return lambda x: -x
+    else:
+        try:
+            return getattr(F, name)
+        except AttributeError:
+            raise ValueError(f"Unknown activation function: {name}")
+
+
+def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any:
+    if chunk_size <= 0:
+        return func(*args, **kwargs)
+    B = None
+    for arg in list(args) + list(kwargs.values()):
+        if isinstance(arg, torch.Tensor):
+            B = arg.shape[0]
+            break
+    assert (
+        B is not None
+    ), "No tensor found in args or kwargs, cannot determine batch size."
+    out = defaultdict(list)
+    out_type = None
+    # max(1, B) to support B == 0
+    for i in range(0, max(1, B), chunk_size):
+        out_chunk = func(
+            *[
+                arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for arg in args
+            ],
+            **{
+                k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for k, arg in kwargs.items()
+            },
+        )
+        if out_chunk is None:
+            continue
+        out_type = type(out_chunk)
+        if isinstance(out_chunk, torch.Tensor):
+            out_chunk = {0: out_chunk}
+        elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list):
+            chunk_length = len(out_chunk)
+            out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)}
+        elif isinstance(out_chunk, dict):
+            pass
+        else:
+            print(
+                f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}."
+            )
+            exit(1)
+        for k, v in out_chunk.items():
+            v = v if torch.is_grad_enabled() else v.detach()
+            out[k].append(v)
+
+    if out_type is None:
+        return None
+
+    out_merged: Dict[Any, Optional[torch.Tensor]] = {}
+    for k, v in out.items():
+        if all([vv is None for vv in v]):
+            # allow None in return value
+            out_merged[k] = None
+        elif all([isinstance(vv, torch.Tensor) for vv in v]):
+            out_merged[k] = torch.cat(v, dim=0)
+        else:
+            raise TypeError(
+                f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}"
+            )
+
+    if out_type is torch.Tensor:
+        return out_merged[0]
+    elif out_type in [tuple, list]:
+        return out_type([out_merged[i] for i in range(chunk_length)])
+    elif out_type is dict:
+        return out_merged
+
+
+def get_ray_directions(
+    H: int,
+    W: int,
+    focal: Union[float, Tuple[float, float]],
+    principal: Optional[Tuple[float, float]] = None,
+    use_pixel_centers: bool = True,
+    normalize: bool = True,
+) -> Float[Tensor, "H W 3"]:
+    """
+    Get ray directions for all pixels in camera coordinate.
+    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
+               ray-tracing-generating-camera-rays/standard-coordinate-systems
+
+    Inputs:
+        H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
+    Outputs:
+        directions: (H, W, 3), the direction of the rays in camera coordinate
+    """
+    pixel_center = 0.5 if use_pixel_centers else 0
+
+    if isinstance(focal, float):
+        fx, fy = focal, focal
+        cx, cy = W / 2, H / 2
+    else:
+        fx, fy = focal
+        assert principal is not None
+        cx, cy = principal
+
+    i, j = torch.meshgrid(
+        torch.arange(W, dtype=torch.float32) + pixel_center,
+        torch.arange(H, dtype=torch.float32) + pixel_center,
+        indexing="xy",
+    )
+
+    directions: Float[Tensor, "H W 3"] = torch.stack(
+        [(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1
+    )
+
+    if normalize:
+        directions = F.normalize(directions, dim=-1)
+
+    return directions
+
+
+def get_rays(
+    directions: Float[Tensor, "... 3"],
+    c2w: Float[Tensor, "... 4 4"],
+    keepdim=False,
+    noise_scale=0.0,
+    normalize=False,
+) -> Tuple[Float[Tensor, "... 3"], Float[Tensor, "... 3"]]:
+    # Rotate ray directions from camera coordinate to the world coordinate
+    assert directions.shape[-1] == 3
+
+    if directions.ndim == 2:  # (N_rays, 3)
+        if c2w.ndim == 2:  # (4, 4)
+            c2w = c2w[None, :, :]
+        assert c2w.ndim == 3  # (N_rays, 4, 4) or (1, 4, 4)
+        rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1)  # (N_rays, 3)
+        rays_o = c2w[:, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 3:  # (H, W, 3)
+        assert c2w.ndim in [2, 3]
+        if c2w.ndim == 2:  # (4, 4)
+            rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
+                -1
+            )  # (H, W, 3)
+            rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
+        elif c2w.ndim == 3:  # (B, 4, 4)
+            rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+                -1
+            )  # (B, H, W, 3)
+            rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 4:  # (B, H, W, 3)
+        assert c2w.ndim == 3  # (B, 4, 4)
+        rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+            -1
+        )  # (B, H, W, 3)
+        rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+
+    # add camera noise to avoid grid-like artifect
+    # https://github.com/ashawkey/stable-dreamfusion/blob/49c3d4fa01d68a4f027755acf94e1ff6020458cc/nerf/utils.py#L373
+    if noise_scale > 0:
+        rays_o = rays_o + torch.randn(3, device=rays_o.device) * noise_scale
+        rays_d = rays_d + torch.randn(3, device=rays_d.device) * noise_scale
+
+    if normalize:
+        rays_d = F.normalize(rays_d, dim=-1)
+    if not keepdim:
+        rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
+
+    return rays_o, rays_d
+
+
+def get_projection_matrix(
+    fovy: Union[float, Float[Tensor, "B"]], aspect_wh: float, near: float, far: float
+) -> Float[Tensor, "*B 4 4"]:
+    if isinstance(fovy, float):
+        proj_mtx = torch.zeros(4, 4, dtype=torch.float32)
+        proj_mtx[0, 0] = 1.0 / (math.tan(fovy / 2.0) * aspect_wh)
+        proj_mtx[1, 1] = -1.0 / math.tan(
+            fovy / 2.0
+        )  # add a negative sign here as the y axis is flipped in nvdiffrast output
+        proj_mtx[2, 2] = -(far + near) / (far - near)
+        proj_mtx[2, 3] = -2.0 * far * near / (far - near)
+        proj_mtx[3, 2] = -1.0
+    else:
+        batch_size = fovy.shape[0]
+        proj_mtx = torch.zeros(batch_size, 4, 4, dtype=torch.float32)
+        proj_mtx[:, 0, 0] = 1.0 / (torch.tan(fovy / 2.0) * aspect_wh)
+        proj_mtx[:, 1, 1] = -1.0 / torch.tan(
+            fovy / 2.0
+        )  # add a negative sign here as the y axis is flipped in nvdiffrast output
+        proj_mtx[:, 2, 2] = -(far + near) / (far - near)
+        proj_mtx[:, 2, 3] = -2.0 * far * near / (far - near)
+        proj_mtx[:, 3, 2] = -1.0
+    return proj_mtx
+
+
+def get_mvp_matrix(
+    c2w: Float[Tensor, "*B 4 4"], proj_mtx: Float[Tensor, "*B 4 4"]
+) -> Float[Tensor, "*B 4 4"]:
+    # calculate w2c from c2w: R' = Rt, t' = -Rt * t
+    # mathematically equivalent to (c2w)^-1
+    if c2w.ndim == 2:
+        assert proj_mtx.ndim == 2
+        w2c: Float[Tensor, "4 4"] = torch.zeros(4, 4).to(c2w)
+        w2c[:3, :3] = c2w[:3, :3].permute(1, 0)
+        w2c[:3, 3:] = -c2w[:3, :3].permute(1, 0) @ c2w[:3, 3:]
+        w2c[3, 3] = 1.0
+    else:
+        w2c: Float[Tensor, "B 4 4"] = torch.zeros(c2w.shape[0], 4, 4).to(c2w)
+        w2c[:, :3, :3] = c2w[:, :3, :3].permute(0, 2, 1)
+        w2c[:, :3, 3:] = -c2w[:, :3, :3].permute(0, 2, 1) @ c2w[:, :3, 3:]
+        w2c[:, 3, 3] = 1.0
+    # calculate mvp matrix by proj_mtx @ w2c (mv_mtx)
+    mvp_mtx = proj_mtx @ w2c
+    return mvp_mtx
+
+
+def get_intrinsic_from_fov(fov, H, W, bs=-1):
+    focal_length = 0.5 * H / np.tan(0.5 * fov)
+    intrinsic = np.identity(3, dtype=np.float32)
+    intrinsic[0, 0] = focal_length
+    intrinsic[1, 1] = focal_length
+    intrinsic[0, 2] = W / 2.0
+    intrinsic[1, 2] = H / 2.0
+
+    if bs > 0:
+        intrinsic = intrinsic[None].repeat(bs, axis=0)
+
+    return torch.from_numpy(intrinsic)
+
+
+def binary_cross_entropy(input, target):
+    """
+    F.binary_cross_entropy is not numerically stable in mixed-precision training.
+    """
+    return -(target * torch.log(input) + (1 - target) * torch.log(1 - input)).mean()
+
+
+def tet_sdf_diff(
+    vert_sdf: Float[Tensor, "Nv 1"], tet_edges: Integer[Tensor, "Ne 2"]
+) -> Float[Tensor, ""]:
+    sdf_f1x6x2 = vert_sdf[:, 0][tet_edges.reshape(-1)].reshape(-1, 2)
+    mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1])
+    sdf_f1x6x2 = sdf_f1x6x2[mask]
+    sdf_diff = F.binary_cross_entropy_with_logits(
+        sdf_f1x6x2[..., 0], (sdf_f1x6x2[..., 1] > 0).float()
+    ) + F.binary_cross_entropy_with_logits(
+        sdf_f1x6x2[..., 1], (sdf_f1x6x2[..., 0] > 0).float()
+    )
+    return sdf_diff
+
+
+def validate_empty_rays(ray_indices, t_start, t_end):
+    if ray_indices.nelement() == 0:
+        warn("Empty rays_indices!")
+        ray_indices = torch.LongTensor([0]).to(ray_indices)
+        t_start = torch.Tensor([0]).to(ray_indices)
+        t_end = torch.Tensor([0]).to(ray_indices)
+    return ray_indices, t_start, t_end
+
+
+def rays_intersect_bbox(
+    rays_o: Float[Tensor, "N 3"],
+    rays_d: Float[Tensor, "N 3"],
+    radius: Float,
+    near: Float = 0.0,
+    valid_thresh: Float = 0.01,
+    background: bool = False,
+):
+    input_shape = rays_o.shape[:-1]
+    rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3)
+    rays_d_valid = torch.where(
+        rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d
+    )
+    if type(radius) in [int, float]:
+        radius = torch.FloatTensor(
+            [[-radius, radius], [-radius, radius], [-radius, radius]]
+        ).to(rays_o.device)
+    radius = (
+        1.0 - 1.0e-3
+    ) * radius  # tighten the radius to make sure the intersection point lies in the bounding box
+    interx0 = (radius[..., 1] - rays_o) / rays_d_valid
+    interx1 = (radius[..., 0] - rays_o) / rays_d_valid
+    t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near)
+    t_far = torch.maximum(interx0, interx1).amin(dim=-1)
+
+    # check wheter a ray intersects the bbox or not
+    rays_valid = t_far - t_near > valid_thresh
+
+    t_near_valid, t_far_valid = t_near[rays_valid], t_far[rays_valid]
+    global_near = t_near_valid.min().item()
+    global_far = t_far_valid.max().item()
+
+    t_near[torch.where(~rays_valid)] = 0.0
+    t_far[torch.where(~rays_valid)] = 0.0
+
+    t_near = t_near.view(*input_shape, 1)
+    t_far = t_far.view(*input_shape, 1)
+    rays_valid = rays_valid.view(*input_shape)
+
+    return t_near, t_far, rays_valid
+
+
+def get_plucker_rays(
+    rays_o: Float[Tensor, "*N 3"], rays_d: Float[Tensor, "*N 3"]
+) -> Float[Tensor, "*N 6"]:
+    rays_o = F.normalize(rays_o, dim=-1)
+    rays_d = F.normalize(rays_d, dim=-1)
+    return torch.cat([rays_o.cross(rays_d), rays_d], dim=-1)
+
+
+def c2w_to_polar(c2w: Float[Tensor, "4 4"]) -> Tuple[float, float, float]:
+    cam_pos = c2w[:3, 3]
+    x, y, z = cam_pos.tolist()
+    distance = cam_pos.norm().item()
+    elevation = math.asin(z / distance)
+    if abs(x) < 1.0e-5 and abs(y) < 1.0e-5:
+        azimuth = 0
+    else:
+        azimuth = math.atan2(y, x)
+        if azimuth < 0:
+            azimuth += 2 * math.pi
+
+    return elevation, azimuth, distance
+
+
+def polar_to_c2w(
+    elevation: float, azimuth: float, distance: float
+) -> Float[Tensor, "4 4"]:
+    """
+    Compute L = p - C.
+    Normalize L.
+    Compute s = L x u. (cross product)
+    Normalize s.
+    Compute u' = s x L.
+    rotation = [s, u, -l]
+    """
+    z = distance * math.sin(elevation)
+    x = distance * math.cos(elevation) * math.cos(azimuth)
+    y = distance * math.cos(elevation) * math.sin(azimuth)
+    l = -torch.as_tensor([x, y, z]).float()
+    l = F.normalize(l, dim=0)
+    u = torch.as_tensor([0.0, 0.0, 1.0]).float()
+    s = l.cross(u)
+    s = F.normalize(s, dim=0)
+    u = s.cross(l)
+    rot = torch.stack([s, u, -l], dim=0).T
+    c2w = torch.zeros((4, 4), dtype=torch.float32)
+    c2w[:3, :3] = rot
+    c2w[:3, 3] = torch.as_tensor([x, y, z])
+    c2w[3, 3] = 1
+    return c2w
+
+
+def fourier_position_encoding(x, n_freq: int, dim: int):
+    assert n_freq > 0
+    input_shape = x.shape
+    input_ndim = x.ndim
+    if dim < 0:
+        dim = input_ndim + dim
+    bands = 2 ** torch.arange(n_freq, dtype=x.dtype, device=x.device)
+    for i in range(dim + 1):
+        bands = bands.unsqueeze(0)
+    for i in range(input_ndim - dim - 1):
+        bands = bands.unsqueeze(-1)
+    x = x.view(*input_shape[: dim + 1], 1, *input_shape[dim + 1 :])
+    x = torch.cat(
+        [
+            torch.sin(bands * x).reshape(
+                *input_shape[:dim], -1, *input_shape[dim + 1 :]
+            ),
+            torch.cos(bands * x).reshape(
+                *input_shape[:dim], -1, *input_shape[dim + 1 :]
+            ),
+        ],
+        dim=dim,
+    )
+    return x

Gen_3D_Modules/MV_Adapter/mvadapter/utils/saving.py

@@ -0,0 +1,544 @@
+import json
+import math
+import os
+import re
+import shutil
+from typing import List, Optional, Union
+
+import cv2
+import imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+# import wandb
+from matplotlib import cm
+from matplotlib.colors import LinearSegmentedColormap
+from PIL import Image, ImageDraw
+
+from .typing import *
+
+
+def image_to_tensor(
+    images: List[Image.Image], device: str = "cuda"
+) -> torch.FloatTensor:
+    tensors = []
+    for image in images:
+        tensor = torch.from_numpy(np.array(image)).float() / 255.0
+        tensor = tensor.to(device)
+        tensors.append(tensor)
+    return torch.stack(tensors)
+
+
+def tensor_to_image(
+    data: Union[Image.Image, torch.Tensor, np.ndarray],
+    batched: bool = False,
+    format: str = "HWC",
+) -> Union[Image.Image, List[Image.Image]]:
+    if isinstance(data, Image.Image):
+        return data
+    if isinstance(data, torch.Tensor):
+        data = data.detach().cpu().numpy()
+    if data.dtype == np.float32 or data.dtype == np.float16:
+        data = (data * 255).astype(np.uint8)
+    elif data.dtype == np.bool_:
+        data = data.astype(np.uint8) * 255
+    assert data.dtype == np.uint8
+    if format == "CHW":
+        if batched and data.ndim == 4:
+            data = data.transpose((0, 2, 3, 1))
+        elif not batched and data.ndim == 3:
+            data = data.transpose((1, 2, 0))
+
+    if batched:
+        return [Image.fromarray(d) for d in data]
+    return Image.fromarray(data)
+
+
+def largest_factor_near_sqrt(n: int) -> int:
+    """
+    Finds the largest factor of n that is closest to the square root of n.
+
+    Args:
+        n (int): The integer for which to find the largest factor near its square root.
+
+    Returns:
+        int: The largest factor of n that is closest to the square root of n.
+    """
+    sqrt_n = int(math.sqrt(n))  # Get the integer part of the square root
+
+    # First, check if the square root itself is a factor
+    if sqrt_n * sqrt_n == n:
+        return sqrt_n
+
+    # Otherwise, find the largest factor by iterating from sqrt_n downwards
+    for i in range(sqrt_n, 0, -1):
+        if n % i == 0:
+            return i
+
+    # If n is 1, return 1
+    return 1
+
+
+def make_image_grid(
+    images: List[Image.Image],
+    rows: Optional[int] = None,
+    cols: Optional[int] = None,
+    resize: Optional[int] = None,
+) -> Image.Image:
+    """
+    Prepares a single grid of images. Useful for visualization purposes.
+    """
+    if rows is None and cols is not None:
+        assert len(images) % cols == 0
+        rows = len(images) // cols
+    elif cols is None and rows is not None:
+        assert len(images) % rows == 0
+        cols = len(images) // rows
+    elif rows is None and cols is None:
+        rows = largest_factor_near_sqrt(len(images))
+        cols = len(images) // rows
+
+    assert len(images) == rows * cols
+
+    if resize is not None:
+        images = [img.resize((resize, resize)) for img in images]
+
+    w, h = images[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+
+    for i, img in enumerate(images):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+class SaverMixin:
+    _save_dir: Optional[str] = None
+    _wandb_logger: Optional[Any] = None
+
+    def set_save_dir(self, save_dir: str):
+        self._save_dir = save_dir
+
+    def get_save_dir(self):
+        if self._save_dir is None:
+            raise ValueError("Save dir is not set")
+        return self._save_dir
+
+    def convert_data(self, data):
+        if data is None:
+            return None
+        elif isinstance(data, np.ndarray):
+            return data
+        elif isinstance(data, torch.Tensor):
+            if data.dtype in [torch.float16, torch.bfloat16]:
+                data = data.float()
+            return data.detach().cpu().numpy()
+        elif isinstance(data, list):
+            return [self.convert_data(d) for d in data]
+        elif isinstance(data, dict):
+            return {k: self.convert_data(v) for k, v in data.items()}
+        else:
+            raise TypeError(
+                "Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
+                type(data),
+            )
+
+    def get_save_path(self, filename):
+        save_path = os.path.join(self.get_save_dir(), filename)
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        return save_path
+
+    DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
+    DEFAULT_UV_KWARGS = {
+        "data_format": "HWC",
+        "data_range": (0, 1),
+        "cmap": "checkerboard",
+    }
+    DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"}
+    DEFAULT_GRID_KWARGS = {"align": "max"}
+
+    def get_rgb_image_(self, img, data_format, data_range, rgba=False):
+        img = self.convert_data(img)
+        assert data_format in ["CHW", "HWC"]
+        if data_format == "CHW":
+            img = img.transpose(1, 2, 0)
+        if img.dtype != np.uint8:
+            img = img.clip(min=data_range[0], max=data_range[1])
+            img = (
+                (img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0
+            ).astype(np.uint8)
+        nc = 4 if rgba else 3
+        imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)]
+        imgs = [
+            (
+                img_
+                if img_.shape[-1] == nc
+                else np.concatenate(
+                    [
+                        img_,
+                        np.zeros(
+                            (img_.shape[0], img_.shape[1], nc - img_.shape[2]),
+                            dtype=img_.dtype,
+                        ),
+                    ],
+                    axis=-1,
+                )
+            )
+            for img_ in imgs
+        ]
+        img = np.concatenate(imgs, axis=1)
+        if rgba:
+            img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
+        else:
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        return img
+
+    def _save_rgb_image(
+        self,
+        filename,
+        img,
+        data_format,
+        data_range,
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ):
+        img = self.get_rgb_image_(img, data_format, data_range)
+        cv2.imwrite(filename, img)
+        if name and self._wandb_logger:
+            self._wandb_logger.log_image(
+                key=name, images=[self.get_save_path(filename)], step=step
+            )
+
+    def save_rgb_image(
+        self,
+        filename,
+        img,
+        data_format=DEFAULT_RGB_KWARGS["data_format"],
+        data_range=DEFAULT_RGB_KWARGS["data_range"],
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ) -> str:
+        save_path = self.get_save_path(filename)
+        self._save_rgb_image(save_path, img, data_format, data_range, name, step)
+        return save_path
+
+    def get_uv_image_(self, img, data_format, data_range, cmap):
+        img = self.convert_data(img)
+        assert data_format in ["CHW", "HWC"]
+        if data_format == "CHW":
+            img = img.transpose(1, 2, 0)
+        img = img.clip(min=data_range[0], max=data_range[1])
+        img = (img - data_range[0]) / (data_range[1] - data_range[0])
+        assert cmap in ["checkerboard", "color"]
+        if cmap == "checkerboard":
+            n_grid = 64
+            mask = (img * n_grid).astype(int)
+            mask = (mask[..., 0] + mask[..., 1]) % 2 == 0
+            img = np.ones((img.shape[0], img.shape[1], 3), dtype=np.uint8) * 255
+            img[mask] = np.array([255, 0, 255], dtype=np.uint8)
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        elif cmap == "color":
+            img_ = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
+            img_[..., 0] = (img[..., 0] * 255).astype(np.uint8)
+            img_[..., 1] = (img[..., 1] * 255).astype(np.uint8)
+            img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR)
+            img = img_
+        return img
+
+    def save_uv_image(
+        self,
+        filename,
+        img,
+        data_format=DEFAULT_UV_KWARGS["data_format"],
+        data_range=DEFAULT_UV_KWARGS["data_range"],
+        cmap=DEFAULT_UV_KWARGS["cmap"],
+    ) -> str:
+        save_path = self.get_save_path(filename)
+        img = self.get_uv_image_(img, data_format, data_range, cmap)
+        cv2.imwrite(save_path, img)
+        return save_path
+
+    def get_grayscale_image_(self, img, data_range, cmap):
+        img = self.convert_data(img)
+        img = np.nan_to_num(img)
+        if data_range is None:
+            img = (img - img.min()) / (img.max() - img.min())
+        else:
+            img = img.clip(data_range[0], data_range[1])
+            img = (img - data_range[0]) / (data_range[1] - data_range[0])
+        assert cmap in [None, "jet", "magma", "spectral"]
+        if cmap == None:
+            img = (img * 255.0).astype(np.uint8)
+            img = np.repeat(img[..., None], 3, axis=2)
+        elif cmap == "jet":
+            img = (img * 255.0).astype(np.uint8)
+            img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
+        elif cmap == "magma":
+            img = 1.0 - img
+            base = cm.get_cmap("magma")
+            num_bins = 256
+            colormap = LinearSegmentedColormap.from_list(
+                f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins
+            )(np.linspace(0, 1, num_bins))[:, :3]
+            a = np.floor(img * 255.0)
+            b = (a + 1).clip(max=255.0)
+            f = img * 255.0 - a
+            a = a.astype(np.uint16).clip(0, 255)
+            b = b.astype(np.uint16).clip(0, 255)
+            img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None]
+            img = (img * 255.0).astype(np.uint8)
+        elif cmap == "spectral":
+            colormap = plt.get_cmap("Spectral")
+
+            def blend_rgba(image):
+                image = image[..., :3] * image[..., -1:] + (
+                    1.0 - image[..., -1:]
+                )  # blend A to RGB
+                return image
+
+            img = colormap(img)
+            img = blend_rgba(img)
+            img = (img * 255).astype(np.uint8)
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        return img
+
+    def _save_grayscale_image(
+        self,
+        filename,
+        img,
+        data_range,
+        cmap,
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ):
+        img = self.get_grayscale_image_(img, data_range, cmap)
+        cv2.imwrite(filename, img)
+        if name and self._wandb_logger:
+            self._wandb_logger.log_image(
+                key=name, images=[self.get_save_path(filename)], step=step
+            )
+
+    def save_grayscale_image(
+        self,
+        filename,
+        img,
+        data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
+        cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ) -> str:
+        save_path = self.get_save_path(filename)
+        self._save_grayscale_image(save_path, img, data_range, cmap, name, step)
+        return save_path
+
+    def get_image_grid_(self, imgs, align):
+        if isinstance(imgs[0], list):
+            return np.concatenate(
+                [self.get_image_grid_(row, align) for row in imgs], axis=0
+            )
+        cols = []
+        for col in imgs:
+            assert col["type"] in ["rgb", "uv", "grayscale"]
+            if col["type"] == "rgb":
+                rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy()
+                rgb_kwargs.update(col["kwargs"])
+                cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs))
+            elif col["type"] == "uv":
+                uv_kwargs = self.DEFAULT_UV_KWARGS.copy()
+                uv_kwargs.update(col["kwargs"])
+                cols.append(self.get_uv_image_(col["img"], **uv_kwargs))
+            elif col["type"] == "grayscale":
+                grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy()
+                grayscale_kwargs.update(col["kwargs"])
+                cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs))
+
+        if align == "max":
+            h = max([col.shape[0] for col in cols])
+        elif align == "min":
+            h = min([col.shape[0] for col in cols])
+        elif isinstance(align, int):
+            h = align
+        else:
+            raise ValueError(
+                f"Unsupported image grid align: {align}, should be min, max, or int"
+            )
+
+        for i in range(len(cols)):
+            if cols[i].shape[0] != h:
+                w = int(cols[i].shape[1] * h / cols[i].shape[0])
+                cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_CUBIC)
+        return np.concatenate(cols, axis=1)
+
+    def save_image_grid(
+        self,
+        filename,
+        imgs,
+        align=DEFAULT_GRID_KWARGS["align"],
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+        texts: Optional[List[float]] = None,
+    ):
+        save_path = self.get_save_path(filename)
+        img = self.get_image_grid_(imgs, align=align)
+
+        if texts is not None:
+            img = Image.fromarray(img)
+            draw = ImageDraw.Draw(img)
+            black, white = (0, 0, 0), (255, 255, 255)
+            for i, text in enumerate(texts):
+                draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
+                draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
+                draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
+                draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
+                draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
+            img = np.asarray(img)
+
+        cv2.imwrite(save_path, img)
+        if name and self._wandb_logger:
+            self._wandb_logger.log_image(key=name, images=[save_path], step=step)
+        return save_path
+
+    def save_image(self, filename, img) -> str:
+        save_path = self.get_save_path(filename)
+        img = self.convert_data(img)
+        assert img.dtype == np.uint8 or img.dtype == np.uint16
+        if img.ndim == 3 and img.shape[-1] == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        elif img.ndim == 3 and img.shape[-1] == 4:
+            img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
+        cv2.imwrite(save_path, img)
+        return save_path
+
+    def save_cubemap(self, filename, img, data_range=(0, 1), rgba=False) -> str:
+        save_path = self.get_save_path(filename)
+        img = self.convert_data(img)
+        assert img.ndim == 4 and img.shape[0] == 6 and img.shape[1] == img.shape[2]
+
+        imgs_full = []
+        for start in range(0, img.shape[-1], 3):
+            img_ = img[..., start : start + 3]
+            img_ = np.stack(
+                [
+                    self.get_rgb_image_(img_[i], "HWC", data_range, rgba=rgba)
+                    for i in range(img_.shape[0])
+                ],
+                axis=0,
+            )
+            size = img_.shape[1]
+            placeholder = np.zeros((size, size, 3), dtype=np.float32)
+            img_full = np.concatenate(
+                [
+                    np.concatenate(
+                        [placeholder, img_[2], placeholder, placeholder], axis=1
+                    ),
+                    np.concatenate([img_[1], img_[4], img_[0], img_[5]], axis=1),
+                    np.concatenate(
+                        [placeholder, img_[3], placeholder, placeholder], axis=1
+                    ),
+                ],
+                axis=0,
+            )
+            imgs_full.append(img_full)
+
+        imgs_full = np.concatenate(imgs_full, axis=1)
+        cv2.imwrite(save_path, imgs_full)
+        return save_path
+
+    def save_data(self, filename, data) -> str:
+        data = self.convert_data(data)
+        if isinstance(data, dict):
+            if not filename.endswith(".npz"):
+                filename += ".npz"
+            save_path = self.get_save_path(filename)
+            np.savez(save_path, **data)
+        else:
+            if not filename.endswith(".npy"):
+                filename += ".npy"
+            save_path = self.get_save_path(filename)
+            np.save(save_path, data)
+        return save_path
+
+    def save_state_dict(self, filename, data) -> str:
+        save_path = self.get_save_path(filename)
+        torch.save(data, save_path)
+        return save_path
+
+    def save_img_sequence(
+        self,
+        filename,
+        img_dir,
+        matcher,
+        save_format="mp4",
+        fps=30,
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ) -> str:
+        assert save_format in ["gif", "mp4"]
+        if not filename.endswith(save_format):
+            filename += f".{save_format}"
+        save_path = self.get_save_path(filename)
+        matcher = re.compile(matcher)
+        img_dir = os.path.join(self.get_save_dir(), img_dir)
+        imgs = []
+        for f in os.listdir(img_dir):
+            if matcher.search(f):
+                imgs.append(f)
+        imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
+        imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
+
+        if save_format == "gif":
+            imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
+            imageio.mimsave(save_path, imgs, fps=fps, palettesize=256)
+        elif save_format == "mp4":
+            imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
+            imageio.mimsave(save_path, imgs, fps=fps)
+        if name and self._wandb_logger:
+            from .core import warn
+
+            warn("Wandb logger does not support video logging yet!")
+        return save_path
+
+    def save_img_sequences(
+        self,
+        seq_dir,
+        matcher,
+        save_format="mp4",
+        fps=30,
+        delete=True,
+        name: Optional[str] = None,
+        step: Optional[int] = None,
+    ):
+        seq_dir_ = os.path.join(self.get_save_dir(), seq_dir)
+        for f in os.listdir(seq_dir_):
+            img_dir_ = os.path.join(seq_dir_, f)
+            if not os.path.isdir(img_dir_):
+                continue
+            try:
+                self.save_img_sequence(
+                    os.path.join(seq_dir, f),
+                    os.path.join(seq_dir, f),
+                    matcher,
+                    save_format=save_format,
+                    fps=fps,
+                    name=f"{name}_{f}",
+                    step=step,
+                )
+                if delete:
+                    shutil.rmtree(img_dir_)
+            except:
+                from .core import warn
+
+                warn(f"Video saving for directory {seq_dir_} failed!")
+
+    def save_file(self, filename, src_path, delete=False) -> str:
+        save_path = self.get_save_path(filename)
+        shutil.copyfile(src_path, save_path)
+        if delete:
+            os.remove(src_path)
+        return save_path
+
+    def save_json(self, filename, payload) -> str:
+        save_path = self.get_save_path(filename)
+        with open(save_path, "w") as f:
+            f.write(json.dumps(payload))
+        return save_path

Gen_3D_Modules/MV_Adapter/mvadapter/utils/typing.py

@@ -0,0 +1,40 @@
+"""
+This module contains type annotations for the project, using
+1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects
+2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors
+
+Two types of typing checking can be used:
+1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode)
+2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking)
+"""
+
+# Basic types
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    NamedTuple,
+    NewType,
+    Optional,
+    Sized,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+
+# Tensor dtype
+# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
+from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
+
+# Config type
+from omegaconf import DictConfig, ListConfig
+
+# PyTorch Tensor type
+from torch import Tensor
+
+# Runtime type checking decorator
+from typeguard import typechecked as typechecker

Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py

@@ -0,0 +1,672 @@
+"""
+Utilities for ComfyUI MV-Adapter nodes
+"""
+
+import os
+import torch
+import numpy as np
+from PIL import Image
+from typing import Optional, Union, List
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
+from huggingface_hub import snapshot_download
+
+from .mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
+from .mvadapter.pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
+from .mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
+from .mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
+from .mvadapter.utils.mesh_utils import (
+    NVDiffRastContextWrapper,
+    get_orthogonal_camera,
+    load_mesh,
+    render,
+)
+from .mvadapter.utils import make_image_grid, tensor_to_image
+from .mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
+
+try:
+    from mmgp import offload, profile_type
+    MMGP_AVAILABLE = True
+except ImportError:
+    MMGP_AVAILABLE = False
+    print("Warning: mmgp not available, memory optimization disabled")
+
+
+def prepare_pipeline(
+    base_model: str,
+    vae_model: Optional[str] = None,
+    unet_model: Optional[str] = None,
+    lora_model: Optional[str] = None,
+    adapter_path: str = "huanngzh/mv-adapter",
+    scheduler: str = "default",
+    num_views: int = 6,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.float16,
+    use_mmgp: bool = False,
+    adapter_local_path: Optional[str] = None,
+):
+    """
+    Prepare MV-Adapter pipeline
+    
+    Args:
+        base_model: Base Stable Diffusion XL model
+        vae_model: Custom VAE model (optional)
+        unet_model: Custom UNet model (optional)
+        lora_model: LoRA model (optional)
+        adapter_path: Path to adapter weights
+        scheduler: Scheduler type ("default", "ddpm", "lcm")
+        num_views: Number of views to generate
+        device: Device for computation
+        dtype: Data type
+        use_mmgp: Use mmgp for memory optimization
+        adapter_local_path: Local path for adapter download
+    """
+    # Load vae and unet if provided
+    pipe_kwargs = {}
+    if vae_model is not None:
+        pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
+    if unet_model is not None:
+        pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
+
+    # Prepare pipeline
+    pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
+
+    # Load scheduler if provided
+    if scheduler != "default":
+        scheduler_class = None
+        if scheduler == "ddpm":
+            scheduler_class = DDPMScheduler
+        elif scheduler == "lcm":
+            scheduler_class = LCMScheduler
+
+        pipe.scheduler = ShiftSNRScheduler.from_scheduler(
+            pipe.scheduler,
+            shift_mode="interpolated",
+            shift_scale=8.0,
+            scheduler_class=scheduler_class,
+        )
+    
+    pipe.init_custom_adapter(
+        num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
+    )
+    
+    # Load adapter weights
+    weight_name = "mvadapter_ig2mv_sdxl.safetensors"
+    pipe.load_custom_adapter(adapter_path, weight_name=weight_name, local_cache_dir=adapter_local_path)
+
+    pipe.to(device=device, dtype=dtype)
+    pipe.cond_encoder.to(device=device, dtype=dtype)
+
+    # load lora if provided
+    if lora_model is not None:
+        model_, name_ = lora_model.rsplit("/", 1)
+        pipe.load_lora_weights(model_, weight_name=name_)
+
+    # Apply mmgp optimization if available and requested
+    if use_mmgp and MMGP_AVAILABLE:
+        import gc
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        all_models = {
+            "vae": pipe.vae,
+            "text_encoder": pipe.text_encoder,
+            "text_encoder_2": pipe.text_encoder_2,
+            "unet": pipe.unet,
+            "cond_encoder": pipe.cond_encoder,  
+        }
+        
+        if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
+            all_models["image_encoder"] = pipe.image_encoder
+        
+        offload.profile(all_models, profile_type.HighRAM_LowVRAM)
+        print("mmgp profiling applied successfully.")
+    
+    # Enable VAE slicing for memory efficiency
+    pipe.enable_vae_slicing()
+
+    return pipe
+
+
+# def create_bg_remover(device: str = "cuda"):
+#     """Create background removal function"""
+#     birefnet = AutoModelForImageSegmentation.from_pretrained(
+#         "ZhengPeng7/BiRefNet", trust_remote_code=True
+#     )
+#     birefnet.to(device)
+#     transform_image = transforms.Compose([
+#         transforms.Resize((1024, 1024)),
+#         transforms.ToTensor(),
+#         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+#     ])
+    
+#     def remove_bg(image):
+#         image_size = image.size
+#         input_images = transform_image(image).unsqueeze(0).to(device)
+#         with torch.no_grad():
+#             preds = birefnet(input_images)[-1].sigmoid().cpu()
+#         pred = preds[0].squeeze()
+#         pred_pil = transforms.ToPILImage()(pred)
+#         mask = pred_pil.resize(image_size)
+#         image.putalpha(mask)
+#         return image
+    
+#     return remove_bg
+
+
+def preprocess_image(image: Image.Image, height: int, width: int) -> Image.Image:
+    """Preprocess image for MV-Adapter"""
+    image = np.array(image)
+    alpha = image[..., 3] > 0
+    H, W = alpha.shape
+    # get the bounding box of alpha
+    y, x = np.where(alpha)
+    y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+    x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
+    image_center = image[y0:y1, x0:x1]
+    # resize the longer side to H * 0.9
+    H, W, _ = image_center.shape
+    if H > W:
+        W = int(W * (height * 0.9) / H)
+        H = int(height * 0.9)
+    else:
+        H = int(H * (width * 0.9) / W)
+        W = int(width * 0.9)
+    image_center = np.array(Image.fromarray(image_center).resize((W, H)))
+    # pad to H, W
+    start_h = (height - H) // 2
+    start_w = (width - W) // 2
+    image = np.zeros((height, width, 4), dtype=np.uint8)
+    image[start_h : start_h + H, start_w : start_w + W] = image_center
+    image = image.astype(np.float32) / 255.0
+    image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+    image = (image * 255).clip(0, 255).astype(np.uint8)
+    image = Image.fromarray(image)
+
+    return image
+
+
+def run_pipeline(
+    pipe,
+    mesh_path: str,
+    num_views: int,
+    text: str,
+    image: Union[str, Image.Image],
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    seed: int,
+    remove_bg_fn=None,
+    reference_conditioning_scale: float = 1.0,
+    negative_prompt: str = "watermark, ugly, deformed, noisy, blurry, low contrast",
+    lora_scale: float = 1.0,
+    device: str = "cuda",
+):
+    """
+    Execute multi-view image generation
+    """
+    # Prepare cameras
+    cameras = get_orthogonal_camera(
+        elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
+        distance=[1.8] * num_views,
+        left=-0.55,
+        right=0.55,
+        bottom=-0.55,
+        top=0.55,
+        azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
+        device=device,
+    )
+    ctx = NVDiffRastContextWrapper(device=device)
+
+    mesh = load_mesh(mesh_path, rescale=True, device=device)
+    render_out = render(
+        ctx,
+        mesh,
+        cameras,
+        height=height,
+        width=width,
+        render_attr=False,
+        normal_background=0.0,
+    )
+    
+    pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
+    normal_images = tensor_to_image(
+        (render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
+    )
+    
+    control_images = (
+        torch.cat(
+            [
+                (render_out.pos + 0.5).clamp(0, 1),
+                (render_out.normal / 2 + 0.5).clamp(0, 1),
+            ],
+            dim=-1,
+        )
+        .permute(0, 3, 1, 2)
+        .to(device)
+    )
+
+    # Prepare image
+    reference_image = Image.open(image) if isinstance(image, str) else image
+    # if remove_bg_fn is not None:
+    #     reference_image = remove_bg_fn(reference_image)
+    #     reference_image = preprocess_image(reference_image, height, width)
+    if reference_image.mode == "RGBA":
+        reference_image = preprocess_image(reference_image, height, width)
+
+    pipe_kwargs = {}
+    if seed != -1 and isinstance(seed, int):
+        pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
+
+    images = pipe(
+        text,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_views,
+        control_image=control_images,
+        control_conditioning_scale=1.0,
+        reference_image=reference_image,
+        reference_conditioning_scale=reference_conditioning_scale,
+        negative_prompt=negative_prompt,
+        cross_attention_kwargs={"scale": lora_scale},
+        **pipe_kwargs,
+    ).images
+
+    return images, pos_images, normal_images, reference_image
+
+
+def create_image_grid(images: List[Image.Image], rows: int = 1) -> Image.Image:
+    return make_image_grid(images, rows=rows)
+
+
+def pil_to_torch_tensor(image: Image.Image) -> torch.Tensor:
+    import numpy as np
+    array = np.array(image)
+    tensor = torch.from_numpy(array).float() / 255.0
+    
+    # Add batch dimension if needed
+    if tensor.dim() == 3:
+        tensor = tensor.unsqueeze(0)
+    
+    return tensor
+
+
+def pils_to_torch_tensor(images: List[Image.Image]) -> torch.Tensor:
+    tensors = [pil_to_torch_tensor(img).squeeze(0) for img in images]
+    return torch.stack(tensors, dim=0)
+
+
+def prepare_t2mv_pipeline(
+    base_model: str,
+    vae_model: Optional[str] = None,
+    unet_model: Optional[str] = None,
+    lora_model: Optional[str] = None,
+    adapter_path: str = "huanngzh/mv-adapter",
+    scheduler: str = "default",
+    num_views: int = 6,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.float16,
+    use_mmgp: bool = False,
+    adapter_local_path: Optional[str] = None,
+):
+    """
+    Prepare MV-Adapter T2MV pipeline
+    
+    Args:
+        base_model: Base Stable Diffusion XL model
+        vae_model: Custom VAE model (optional)
+        unet_model: Custom UNet model (optional)
+        lora_model: LoRA model (optional)
+        adapter_path: Path to adapter weights
+        scheduler: Scheduler type ("default", "ddpm", "lcm")
+        num_views: Number of views to generate
+        device: Device for computation
+        dtype: Data type
+        use_mmgp: Use mmgp for memory optimization
+        adapter_local_path: Local path for adapter download
+    """
+    # Load vae and unet if provided
+    pipe_kwargs = {}
+    if vae_model is not None:
+        pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
+    if unet_model is not None:
+        pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
+
+    # Prepare pipeline
+    pipe = MVAdapterT2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
+
+    # Load scheduler if provided
+    if scheduler != "default":
+        scheduler_class = None
+        if scheduler == "ddpm":
+            scheduler_class = DDPMScheduler
+        elif scheduler == "lcm":
+            scheduler_class = LCMScheduler
+
+        pipe.scheduler = ShiftSNRScheduler.from_scheduler(
+            pipe.scheduler,
+            shift_mode="interpolated",
+            shift_scale=8.0,
+            scheduler_class=scheduler_class,
+        )
+    
+    pipe.init_custom_adapter(num_views=num_views)
+    
+    # Load adapter weights
+    weight_name = "mvadapter_t2mv_sdxl.safetensors"
+    pipe.load_custom_adapter(adapter_path, weight_name=weight_name, local_cache_dir=adapter_local_path)
+
+    pipe.to(device=device, dtype=dtype)
+    pipe.cond_encoder.to(device=device, dtype=dtype)
+
+    # load lora if provided
+    if lora_model is not None:
+        model_, name_ = lora_model.rsplit("/", 1)
+        pipe.load_lora_weights(model_, weight_name=name_)
+
+    # Apply mmgp optimization if available and requested
+    if use_mmgp and MMGP_AVAILABLE:
+        import gc
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        all_models = {
+            "vae": pipe.vae,
+            "text_encoder": pipe.text_encoder,
+            "text_encoder_2": pipe.text_encoder_2,
+            "unet": pipe.unet,
+            "cond_encoder": pipe.cond_encoder,  
+        }
+        
+        if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
+            all_models["image_encoder"] = pipe.image_encoder
+        
+        offload.profile(all_models, profile_type.HighRAM_LowVRAM)
+        print("mmgp profiling applied successfully.")
+    
+    # Enable VAE slicing for memory efficiency
+    pipe.enable_vae_slicing()
+
+    return pipe
+
+
+def run_t2mv_pipeline(
+    pipe,
+    num_views: int,
+    text: str,
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    seed: int,
+    negative_prompt: str = "watermark, ugly, deformed, noisy, blurry, low contrast",
+    lora_scale: float = 1.0,
+    device: str = "cuda",
+    azimuth_deg: Optional[List[int]] = None,
+):
+    """
+    Execute text-to-multiview generation
+    """
+    # Prepare cameras
+    if azimuth_deg is None:
+        azimuth_deg = [0, 45, 90, 180, 270, 315][:num_views]
+    
+    cameras = get_orthogonal_camera(
+        elevation_deg=[0] * num_views,
+        distance=[1.8] * num_views,
+        left=-0.55,
+        right=0.55,
+        bottom=-0.55,
+        top=0.55,
+        azimuth_deg=[x - 90 for x in azimuth_deg],
+        device=device,
+    )
+
+    plucker_embeds = get_plucker_embeds_from_cameras_ortho(
+        cameras.c2w, [1.1] * num_views, width
+    )
+    control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
+
+    pipe_kwargs = {}
+    if seed != -1 and isinstance(seed, int):
+        pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
+
+    images = pipe(
+        text,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_views,
+        control_image=control_images,
+        control_conditioning_scale=1.0,
+        negative_prompt=negative_prompt,
+        cross_attention_kwargs={"scale": lora_scale},
+        **pipe_kwargs,
+    ).images
+
+    return images 
+
+
+def prepare_texture_pipeline(
+    upscaler_ckpt_path: Optional[str] = None,
+    inpaint_ckpt_path: Optional[str] = None,
+    device: str = "cuda",
+    use_mmgp: bool = False,
+):
+    """
+    Prepare texture projection pipeline
+    
+    Args:
+        upscaler_ckpt_path: Path to upscaler checkpoint (optional)
+        inpaint_ckpt_path: Path to inpaint checkpoint (optional)
+        device: Device for computation
+        use_mmgp: Use mmgp for memory optimization
+    """
+    from .mvadapter.pipelines.pipeline_texture import TexturePipeline
+    
+    if upscaler_ckpt_path and not os.path.exists(upscaler_ckpt_path):
+        print(f"Warning: Upscaler checkpoint not found: {upscaler_ckpt_path}, using None")
+        upscaler_ckpt_path = None
+        
+    if inpaint_ckpt_path and not os.path.exists(inpaint_ckpt_path):
+        print(f"Warning: Inpaint checkpoint not found: {inpaint_ckpt_path}, using None")
+        inpaint_ckpt_path = None
+    
+    texture_pipe = TexturePipeline(
+        upscaler_ckpt_path=upscaler_ckpt_path,
+        inpaint_ckpt_path=inpaint_ckpt_path,
+        device=device
+    )
+    
+    # Apply mmgp optimization if available and requested
+    if use_mmgp and MMGP_AVAILABLE:
+        import gc
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        all_models = {}
+        
+        if hasattr(texture_pipe, 'upscaler') and texture_pipe.upscaler is not None:
+            all_models["upscaler"] = texture_pipe.upscaler
+            
+        if hasattr(texture_pipe, 'inpainter') and texture_pipe.inpainter is not None:
+            all_models["inpainter"] = texture_pipe.inpainter
+        
+        if all_models:
+            offload.profile(all_models, profile_type.HighRAM_LowVRAM)
+            print("mmgp profiling applied to texture pipeline successfully.")
+        else:
+            print("No models loaded for texture pipeline, skipping mmgp profiling.")
+    
+    return texture_pipe
+
+
+def run_texture_pipeline(
+    texture_pipe,
+    multiview_images: List[Image.Image],
+    mesh_path: str,
+    save_dir: str,
+    save_name: str = "textured_model",
+    uv_size: int = 4096,
+    view_upscale: bool = True,
+    inpaint_mode: str = "view",
+    uv_unwarp: bool = True,
+    preprocess_mesh: bool = False,
+    move_to_center: bool = False,
+    front_x: bool = True,
+    camera_azimuth_deg: List[float] = None,
+    camera_elevation_deg: List[float] = None,
+    camera_distance: float = 1.0,
+    camera_ortho_scale: float = 1.1,
+    debug_mode: bool = False,
+    apply_dilate: bool = True,
+):
+    """
+    Execute texture projection pipeline
+    
+    Args:
+        texture_pipe: Prepared texture pipeline
+        multiview_images: List of PIL images (6 views)
+        mesh_path: Path to 3D mesh file
+        save_dir: Directory to save results
+        save_name: Name for saved files
+        uv_size: UV texture resolution
+        view_upscale: Enable view upscaling
+        inpaint_mode: Inpainting mode ("none", "uv", "view")
+        uv_unwarp: Enable UV unwrapping
+        preprocess_mesh: Enable mesh preprocessing
+        move_to_center: Move mesh to center
+        front_x: Front face along X axis
+        camera_azimuth_deg: Camera azimuth angles
+        camera_elevation_deg: Camera elevation angles
+        camera_distance: Camera distance
+        camera_ortho_scale: Orthographic scale
+        debug_mode: Enable debug mode
+        apply_dilate: Apply dilate to remove texture seams
+    """
+    from .mvadapter.pipelines.pipeline_texture import ModProcessConfig
+    from .mvadapter.utils import make_image_grid
+    
+    if not mesh_path or not os.path.exists(mesh_path):
+        raise ValueError(f"Mesh file does not exist: {mesh_path}")
+    
+    os.makedirs(save_dir, exist_ok=True)
+    
+    if camera_azimuth_deg is None:
+        camera_azimuth_deg = [0, 90, 180, 270, 180, 180]
+    if camera_elevation_deg is None:
+        camera_elevation_deg = [0, 0, 0, 0, 89.99, -89.99]
+    
+    # Adjust azimuth angles (subtract 90 as in original code)
+    azimuth_deg_corrected = [x - 90 for x in camera_azimuth_deg]
+    
+    grid_image = make_image_grid(multiview_images, rows=1)
+    temp_grid_path = os.path.join(save_dir, f"{save_name}_temp_grid.png")
+    grid_image.save(temp_grid_path)
+    
+    try:
+        rgb_process_config = ModProcessConfig(
+            view_upscale=view_upscale,
+            inpaint_mode=inpaint_mode
+        )
+        
+        output = texture_pipe(
+            mesh_path=mesh_path,
+            save_dir=save_dir,
+            save_name=save_name,
+            move_to_center=move_to_center,
+            front_x=front_x,
+            uv_unwarp=uv_unwarp,
+            preprocess_mesh=preprocess_mesh,
+            uv_size=uv_size,
+            rgb_path=temp_grid_path,
+            rgb_process_config=rgb_process_config,
+            camera_elevation_deg=camera_elevation_deg,
+            camera_azimuth_deg=azimuth_deg_corrected,
+            camera_distance=camera_distance,
+            camera_ortho_scale=camera_ortho_scale,
+            debug_mode=debug_mode,
+            apply_dilate=apply_dilate,
+        )
+        
+        if os.path.exists(temp_grid_path):
+            os.remove(temp_grid_path)
+        
+        return output
+        
+    except Exception as e:
+        if os.path.exists(temp_grid_path):
+            os.remove(temp_grid_path)
+        raise e
+
+
+def download_texture_checkpoints(texture_ckpt_dir, upscaler_ckpt_path, inpaint_ckpt_path):
+    """Download texture checkpoint files if they don't exist"""
+    import subprocess
+    
+    os.makedirs(texture_ckpt_dir, exist_ok=True)
+    
+    checkpoints = [
+        {
+            "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
+            "path": upscaler_ckpt_path,
+            "name": "RealESRGAN upscaler"
+        },
+        {
+            "url": "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", 
+            "path": inpaint_ckpt_path,
+            "name": "Big-LaMa inpainter"
+        }
+    ]
+    
+    for ckpt in checkpoints:
+        if not os.path.exists(ckpt["path"]):
+            print(f"Downloading {ckpt['name']} to {ckpt['path']}...")
+            
+            success = False
+            
+            # Try wget first
+            try:
+                subprocess.run([
+                    "wget", ckpt["url"], "-O", ckpt["path"]
+                ], check=True, capture_output=True, text=True)
+                print(f"Successfully downloaded {ckpt['name']} using wget")
+                success = True
+            except (subprocess.CalledProcessError, FileNotFoundError):
+                pass
+            
+            # Try curl if wget failed
+            if not success:
+                try:
+                    subprocess.run([
+                        "curl", "-L", ckpt["url"], "-o", ckpt["path"]
+                    ], check=True, capture_output=True, text=True)
+                    print(f"Successfully downloaded {ckpt['name']} using curl")
+                    success = True
+                except (subprocess.CalledProcessError, FileNotFoundError):
+                    pass
+            
+            # Try Python urllib as fallback
+            if not success:
+                try:
+                    import urllib.request
+                    urllib.request.urlretrieve(ckpt["url"], ckpt["path"])
+                    print(f"Successfully downloaded {ckpt['name']} using urllib")
+                    success = True
+                except Exception as e:
+                    print(f"Failed to download {ckpt['name']} using urllib: {e}")
+            
+            # If all methods failed
+            if not success:
+                print(f"Failed to download {ckpt['name']}. Please download manually:")
+                print(f"wget {ckpt['url']} -O {ckpt['path']}")
+                print(f"or")
+                print(f"curl -L {ckpt['url']} -o {ckpt['path']}")
+        else:
+            print(f"{ckpt['name']} already exists at {ckpt['path']}") 

Gen_3D_Modules/Unique3D/mesh_reconstruction/render.py

@@ -12,11 +12,11 @@ def _warmup(glctx, device=None):
     tri = tensor([[0, 1, 2]], dtype=torch.int32)
     dr.rasterize(glctx, pos, tri, resolution=[256, 256])
 
-glctx = dr.RasterizeGLContext(output_db=False, device="cuda")
+glctx = dr.RasterizeCudaContext(device="cuda")
 
 class NormalsRenderer:
     
-    _glctx:dr.RasterizeGLContext = None
+    _glctx:dr.RasterizeCudaContext = None
     
     def __init__(
             self,

Gen_3D_Modules/Unique3D/scripts/project_mesh.py

@@ -68,7 +68,7 @@ def _warmup(glctx, device=None):
 
 class Pix2FacesRenderer:
     def __init__(self, device="cuda"):
-        self._glctx = dr.RasterizeGLContext(output_db=False, device=device)
+        self._glctx = dr.RasterizeCudaContext(device=device)
         self.device = device
         _warmup(self._glctx, device)
 

MVs_Algorithms/FlexiCubes/flexicubes_renderer.py

@@ -43,7 +43,7 @@ class FlexiCubesRenderer:
         '''
         v_pos_clip = util.xfm_points(mesh.vertices.unsqueeze(0), mvp)  # Rotate it to camera coordinates
         rast, db = dr.rasterize(
-            dr.RasterizeGLContext(), v_pos_clip, mesh.faces.int(), iter_res)
+            dr.RasterizeCudaContext(), v_pos_clip, mesh.faces.int(), iter_res)
 
         alpha_index = rast[..., -1:] > 0
         inv_alpha_index = rast[..., -1:] <= 0

README.md

@@ -46,6 +46,18 @@ Your ComfyUI Root Directory\python_embeded\python.exe install.py
 - For use cases please check out [Example Workflows](./_Example_Workflows/). [**Last update: 29/April/2025**]
   - **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
 
+- **MV-Adapter**: [huanngzh/MV-Adapter](https://github.com/huanngzh/MV-Adapter)
+  - Two generation methods:
+    - IG2MV: Single image + 3D mesh to multi-view images
+    - T2MV: Text prompt to multi-view images
+    - Texturing: Grid image + 3D mesh to textured mesh
+  - Model weights: https://huggingface.co/huanngzh/mv-adapter
+  - Workflows: [IG2MV](./_Example_Workflows/MV-Adapter/MV-Adapter-ig2mv.json), [T2MV](./_Example_Workflows/MV-Adapter/MV-Adapter-t2mv.json), [Texturing](./_Example_Workflows/MV-Adapter/MV-Adapter-Texturing.json)
+    
+   <video controls autoplay loop src="https://github.com/user-attachments/assets/c03fe162-c2fd-4499-aa8d-43e5863ceb3e"></video>
+   <video controls autoplay loop src="https://github.com/user-attachments/assets/0d67dae6-2167-4bf7-83f6-d7c6c3b8e2e8"></video>
+   <video controls autoplay loop src="https://github.com/user-attachments/assets/943d8d31-30e0-4b9f-968d-a68b5899d83d"></video>
+
 - **Stable3DGen**: [Stable-X/Stable3DGen](https://github.com/Stable-X/Stable3DGen)
   - Two models pipeline:
     - Stable3DGen: Single image to 3D Mesh

_Example_Workflows/MV-Adapter/MV-Adapter-Texturing.json

@@ -0,0 +1,315 @@
+{
+  "id": "4f7b5e42-a56c-480c-8bed-bb1a8a8006cd",
+  "revision": 0,
+  "last_node_id": 26,
+  "last_link_id": 43,
+  "nodes": [
+    {
+      "id": 22,
+      "type": "[Comfy3D] Preview 3DMesh",
+      "pos": [
+        5958.7109375,
+        -101.48223114013672
+      ],
+      "size": [
+        600,
+        500
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mesh_file_path",
+          "type": "STRING",
+          "widget": {
+            "name": "mesh_file_path"
+          },
+          "link": 42
+        }
+      ],
+      "outputs": [],
+      "title": "PBR Mesh",
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Preview 3DMesh"
+      },
+      "widgets_values": [
+        "",
+        null
+      ]
+    },
+    {
+      "id": 17,
+      "type": "[Comfy3D] Preview 3DMesh",
+      "pos": [
+        5346.52099609375,
+        -95.2625503540039
+      ],
+      "size": [
+        600,
+        500
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mesh_file_path",
+          "type": "STRING",
+          "widget": {
+            "name": "mesh_file_path"
+          },
+          "link": 41
+        }
+      ],
+      "outputs": [],
+      "title": "Shaded Mesh",
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Preview 3DMesh"
+      },
+      "widgets_values": [
+        "",
+        null
+      ]
+    },
+    {
+      "id": 24,
+      "type": "[Comfy3D] Load MVAdapter Texture Pipeline",
+      "pos": [
+        5015.30126953125,
+        -68.61071014404297
+      ],
+      "size": [
+        307.4322204589844,
+        130
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "texture_pipeline",
+          "type": "TEXTURE_PIPE",
+          "links": [
+            38
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Load MVAdapter Texture Pipeline"
+      },
+      "widgets_values": [
+        "RealESRGAN_x2plus.pth",
+        "big-lama.pt",
+        false,
+        true
+      ]
+    },
+    {
+      "id": 15,
+      "type": "LoadImage",
+      "pos": [
+        4655.7041015625,
+        132.22010803222656
+      ],
+      "size": [
+        325.982177734375,
+        338.67474365234375
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            39
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "squirel_girl_grid.png",
+        "image"
+      ]
+    },
+    {
+      "id": 26,
+      "type": "[Comfy3D] MVAdapter Texture Projection",
+      "pos": [
+        5015.29931640625,
+        137.44638061523438
+      ],
+      "size": [
+        298.482421875,
+        462
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "texture_pipeline",
+          "type": "TEXTURE_PIPE",
+          "link": 38
+        },
+        {
+          "name": "grid_image",
+          "type": "IMAGE",
+          "link": 39
+        },
+        {
+          "name": "mesh_path",
+          "type": "STRING",
+          "widget": {
+            "name": "mesh_path"
+          },
+          "link": 40
+        }
+      ],
+      "outputs": [
+        {
+          "name": "shaded_model_path",
+          "type": "STRING",
+          "links": [
+            41
+          ]
+        },
+        {
+          "name": "pbr_model_path",
+          "type": "STRING",
+          "links": [
+            42
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] MVAdapter Texture Projection"
+      },
+      "widgets_values": [
+        "",
+        "./output",
+        "textured_model",
+        4096,
+        true,
+        "view",
+        true,
+        false,
+        false,
+        true,
+        true,
+        true,
+        "0,90,180,270,180,180",
+        "0,0,0,0,89.99,-89.99",
+        1,
+        1.1,
+        false
+      ]
+    },
+    {
+      "id": 4,
+      "type": "PrimitiveString",
+      "pos": [
+        4659.12890625,
+        -19.559751510620117
+      ],
+      "size": [
+        319.9439392089844,
+        92.46147918701172
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "STRING",
+          "type": "STRING",
+          "links": [
+            40
+          ]
+        }
+      ],
+      "title": "Input mesh",
+      "properties": {
+        "Node name for S&R": "PrimitiveString"
+      },
+      "widgets_values": [
+        "output/StableGenTest/mesh.glb"
+      ]
+    }
+  ],
+  "links": [
+    [
+      38,
+      24,
+      0,
+      26,
+      0,
+      "TEXTURE_PIPE"
+    ],
+    [
+      39,
+      15,
+      0,
+      26,
+      1,
+      "IMAGE"
+    ],
+    [
+      40,
+      4,
+      0,
+      26,
+      2,
+      "STRING"
+    ],
+    [
+      41,
+      26,
+      0,
+      17,
+      0,
+      "STRING"
+    ],
+    [
+      42,
+      26,
+      1,
+      22,
+      0,
+      "STRING"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 4.321942375150668,
+      "offset": [
+        -4633.552682414396,
+        -263.1857272782662
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "node_versions": {
+      "ComfyUI-3D-Pack": "626a5ee60d6fbe3d76d69cd89261d445dbca2bc4",
+      "comfy-core": "0.3.35"
+    },
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}

_Example_Workflows/MV-Adapter/MV-Adapter-ig2mv.json

@@ -0,0 +1,340 @@
+{
+  "id": "867889cd-eeb4-48b0-853c-21e125a93896",
+  "revision": 0,
+  "last_node_id": 28,
+  "last_link_id": 25,
+  "nodes": [
+    {
+      "id": 8,
+      "type": "PreviewImage",
+      "pos": [
+        3033.075927734375,
+        -927.3829956054688
+      ],
+      "size": [
+        370.7532958984375,
+        288.39501953125
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 6
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 26,
+      "type": "[Comfy3D] Make Image Grid",
+      "pos": [
+        3033.62158203125,
+        -586.97265625
+      ],
+      "size": [
+        270,
+        82
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 22
+        }
+      ],
+      "outputs": [
+        {
+          "name": "image_grid",
+          "type": "IMAGE",
+          "links": [
+            23
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Make Image Grid"
+      },
+      "widgets_values": [
+        1,
+        true
+      ]
+    },
+    {
+      "id": 27,
+      "type": "PreviewImage",
+      "pos": [
+        3032.88818359375,
+        -449.5286865234375
+      ],
+      "size": [
+        420.0350036621094,
+        123.75502014160156
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 23
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 3,
+      "type": "LoadImage",
+      "pos": [
+        2153.30078125,
+        -839.7021484375
+      ],
+      "size": [
+        346.8450622558594,
+        326
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            19
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "squirrel_girl_1.png",
+        "image"
+      ]
+    },
+    {
+      "id": 9,
+      "type": "PrimitiveString",
+      "pos": [
+        2155.323974609375,
+        -452.05352783203125
+      ],
+      "size": [
+        343.5000305175781,
+        67.5549545288086
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "STRING",
+          "type": "STRING",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PrimitiveString"
+      },
+      "widgets_values": [
+        "output/StableGenTest/mesh.glb"
+      ]
+    },
+    {
+      "id": 7,
+      "type": "[Comfy3D] MVAdapter Image To MultiView",
+      "pos": [
+        2557.745361328125,
+        -938.2529296875
+      ],
+      "size": [
+        406.614990234375,
+        407.14501953125
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mvadapter_pipe",
+          "type": "DIFFUSERS_PIPE",
+          "link": 25
+        },
+        {
+          "name": "reference_image",
+          "type": "IMAGE",
+          "link": 19
+        },
+        {
+          "name": "mesh_path",
+          "type": "STRING",
+          "widget": {
+            "name": "mesh_path"
+          },
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "multiview_images",
+          "type": "IMAGE",
+          "links": [
+            6,
+            22
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] MVAdapter Image To MultiView"
+      },
+      "widgets_values": [
+        "",
+        "high quality",
+        "watermark, ugly, deformed, noisy, blurry, low contrast",
+        50,
+        3,
+        1,
+        768,
+        768,
+        710491314281528,
+        "randomize",
+        false,
+        1
+      ]
+    },
+    {
+      "id": 28,
+      "type": "[Comfy3D] Load MVAdapter IG2MV Pipeline",
+      "pos": [
+        2154.555908203125,
+        -1155.862060546875
+      ],
+      "size": [
+        357.2724609375,
+        257.3499755859375
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "mvadapter_pipe",
+          "type": "DIFFUSERS_PIPE",
+          "links": [
+            25
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Load MVAdapter IG2MV Pipeline"
+      },
+      "widgets_values": [
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        "madebyollin/sdxl-vae-fp16-fix",
+        "huanngzh/mv-adapter",
+        "default",
+        6,
+        true,
+        true,
+        "",
+        ""
+      ]
+    }
+  ],
+  "links": [
+    [
+      6,
+      7,
+      0,
+      8,
+      0,
+      "IMAGE"
+    ],
+    [
+      7,
+      9,
+      0,
+      7,
+      2,
+      "STRING"
+    ],
+    [
+      19,
+      3,
+      0,
+      7,
+      1,
+      "IMAGE"
+    ],
+    [
+      22,
+      7,
+      0,
+      26,
+      0,
+      "IMAGE"
+    ],
+    [
+      23,
+      26,
+      0,
+      27,
+      0,
+      "IMAGE"
+    ],
+    [
+      25,
+      28,
+      0,
+      7,
+      0,
+      "DIFFUSERS_PIPE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.9070294784580498,
+      "offset": [
+        -1941.4116551684099,
+        1211.3556830868408
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "node_versions": {
+      "comfy-core": "0.3.35",
+      "ComfyUI-3D-Pack": "626a5ee60d6fbe3d76d69cd89261d445dbca2bc4"
+    },
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}

_Example_Workflows/MV-Adapter/MV-Adapter-t2mv.json

@@ -0,0 +1,160 @@
+{
+  "id": "fd4d59aa-116b-4ec5-920e-f918c36151fe",
+  "revision": 0,
+  "last_node_id": 3,
+  "last_link_id": 2,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": [
+        2673.054931640625,
+        -795.7294311523438
+      ],
+      "size": [
+        291.89996337890625,
+        246
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 1,
+      "type": "[Comfy3D] MVAdapter Text To MultiView",
+      "pos": [
+        2183.7578125,
+        -848.2293701171875
+      ],
+      "size": [
+        400,
+        358
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mvadapter_t2mv_pipe",
+          "type": "DIFFUSERS_PIPE",
+          "link": 2
+        }
+      ],
+      "outputs": [
+        {
+          "name": "multiview_images",
+          "type": "IMAGE",
+          "links": [
+            1
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] MVAdapter Text To MultiView"
+      },
+      "widgets_values": [
+        "Mater, a rusty and beat-up tow truck from the 2006 Disney/Pixar animated film 'Cars', with a rusty brown exterior, big blue eyes.",
+        "watermark, ugly, deformed, noisy, blurry, low contrast",
+        6,
+        50,
+        7,
+        768,
+        768,
+        937383047416961,
+        "randomize",
+        1,
+        "0,45,90,180,270,315"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "[Comfy3D] Load MVAdapter T2MV Pipeline",
+      "pos": [
+        1833.75634765625,
+        -913.3286743164062
+      ],
+      "size": [
+        298.8599548339844,
+        250
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "mvadapter_t2mv_pipe",
+          "type": "DIFFUSERS_PIPE",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "[Comfy3D] Load MVAdapter T2MV Pipeline"
+      },
+      "widgets_values": [
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        "madebyollin/sdxl-vae-fp16-fix",
+        "huanngzh/mv-adapter",
+        "default",
+        6,
+        true,
+        false,
+        "",
+        ""
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      2,
+      3,
+      0,
+      1,
+      0,
+      "DIFFUSERS_PIPE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.2155062499999998,
+      "offset": [
+        -1674.815702398016,
+        990.4808470306099
+      ]
+    },
+    "node_versions": {
+      "comfy-core": "0.3.35",
+      "ComfyUI-3D-Pack": "354ab9f7b56129f1f0fc4a42bea1751b266fd222"
+    },
+    "frontendVersion": "1.20.4",
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}

nodes.py

@@ -103,8 +103,16 @@ from TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline
 from TRELLIS.trellis.utils import postprocessing_utils
 from TripoSG.pipelines.pipeline_triposg import TripoSGPipeline
 from TripoSG.pipelines.pipeline_triposg_scribble import TripoSGScribblePipeline
-from Stable3DGen.trellis.backend_config import get_available_backends, get_available_sparse_backends
 from Stable3DGen.pipeline_builders import StableGenPipelineBuilder
+from MV_Adapter.mvadapter_node_utils import (
+        prepare_pipeline as mvadapter_prepare_pipeline,
+        run_pipeline as mvadapter_run_pipeline, 
+        prepare_t2mv_pipeline as mvadapter_prepare_t2mv_pipeline,
+        run_t2mv_pipeline as mvadapter_run_t2mv_pipeline,
+        prepare_texture_pipeline as mvadapter_prepare_texture_pipeline,
+        download_texture_checkpoints,
+    )
+from MV_Adapter.mvadapter.utils import make_image_grid
 
 os.environ['SPCONV_ALGO'] = 'native'
 
@@ -994,7 +1002,7 @@ class Mesh_Orbit_Renderer:
                 "render_background_color_r": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                 "render_background_color_g": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                 "render_background_color_b": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
-                "force_cuda_rasterize": ("BOOLEAN", {"default": False},),
+                "force_cuda_rasterize": ("BOOLEAN", {"default": True},),
             },
             
             "optional": {
@@ -1309,7 +1317,7 @@ class Fitting_Mesh_With_Multiview_Images:
                 "ms_ssim_loss_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "remesh_after_n_iteration": ("INT", {"default": 512, "min": 128, "max": 100000}),
                 "invert_background_probability": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}),
-                "force_cuda_rasterize": ("BOOLEAN", {"default": False},),
+                "force_cuda_rasterize": ("BOOLEAN", {"default": True},),
             },
         }
 
@@ -4097,34 +4105,33 @@ class Trellis_Structured_3D_Latents_Models:
     ):
         single_image = torch_imgs_to_pils(reference_image, reference_mask)[0]
 
-        with torch.inference_mode(False):
-            outputs = trellis_pipe.run(
-                single_image,
-                # Optional parameters
-                seed=seed,
-                formats=["gaussian", "mesh"],
-                sparse_structure_sampler_params={
-                    "cfg_strength": sparse_structure_guidance_scale,
-                    "steps": sparse_structure_sample_steps,
-                },
-                slat_sampler_params={
-                    "cfg_strength": structured_latent_guidance_scale,
-                    "steps": structured_latent_sample_steps,
-                },
-            )
+        outputs = trellis_pipe.run(
+            single_image,
+            # Optional parameters
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            sparse_structure_sampler_params={
+                "cfg_strength": sparse_structure_guidance_scale,
+                "steps": sparse_structure_sample_steps,
+            },
+            slat_sampler_params={
+                "cfg_strength": structured_latent_guidance_scale,
+                "steps": structured_latent_sample_steps,
+            },
+        )
 
-            # GLB files can be extracted from the outputs
-            vertices, faces, uvs, texture = postprocessing_utils.finalize_mesh(
-                outputs['gaussian'][0],
-                outputs['mesh'][0],
-                # Optional parameters
-                simplify=0.95,          # Ratio of triangles to remove in the simplification process
-                texture_size=1024,      # Size of the texture used for the GLB
-            )
+        # GLB files can be extracted from the outputs
+        vertices, faces, uvs, texture = postprocessing_utils.finalize_mesh(
+            outputs['gaussian'][0],
+            outputs['mesh'][0],
+            # Optional parameters
+            simplify=0.95,          # Ratio of triangles to remove in the simplification process
+            texture_size=1024,      # Size of the texture used for the GLB
+        )
 
-            vertices, faces, uvs, texture = torch.from_numpy(vertices).to(DEVICE), torch.from_numpy(faces).to(torch.int64).to(DEVICE), torch.from_numpy(uvs).to(DEVICE), torch.from_numpy(texture).to(DEVICE)
-            mesh = Mesh(v=vertices, f=faces, vt=uvs, ft=faces, albedo=texture, device=DEVICE)
-            mesh.auto_normal()
+        vertices, faces, uvs, texture = torch.from_numpy(vertices).to(DEVICE), torch.from_numpy(faces).to(torch.int64).to(DEVICE), torch.from_numpy(uvs).to(DEVICE), torch.from_numpy(texture).to(DEVICE)
+        mesh = Mesh(v=vertices, f=faces, vt=uvs, ft=faces, albedo=texture, device=DEVICE)
+        mesh.auto_normal()
 
         return (mesh,)
 
@@ -4237,23 +4244,22 @@ class TripoSG_Scribble_Model:
         
         single_image = torch_imgs_to_pils(scribble_image)[0]
         
-        with torch.inference_mode(False):
-            outputs = tsg_scribble_pipe(
-                image=single_image,
-                prompt=prompt,
-                generator=torch.Generator(device=DEVICE).manual_seed(seed),
-                num_inference_steps=num_inference_steps,
-                guidance_scale=0, # this is a CFG-distilled model
-                attention_kwargs={"cross_attention_scale": prompt_confidence, "cross_attention_2_scale": scribble_confidence},
-                use_flash_decoder=use_flash_decoder,
-                flash_octree_depth=flash_octree_depth, # there're some boundary problems when using flash decoder with this model
-                hierarchical_octree_depth=hierarchical_octree_depth,
-                dense_octree_depth=dense_octree_depth,
-            ).samples[0]
+        outputs = tsg_scribble_pipe(
+            image=single_image,
+            prompt=prompt,
+            generator=torch.Generator(device=DEVICE).manual_seed(seed),
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0, # this is a CFG-distilled model
+            attention_kwargs={"cross_attention_scale": prompt_confidence, "cross_attention_2_scale": scribble_confidence},
+            use_flash_decoder=use_flash_decoder,
+            flash_octree_depth=flash_octree_depth, # there're some boundary problems when using flash decoder with this model
+            hierarchical_octree_depth=hierarchical_octree_depth,
+            dense_octree_depth=dense_octree_depth,
+        ).samples[0]
 
-            vertices, faces = torch.from_numpy(outputs[0].astype(np.float32)).to(DEVICE), torch.from_numpy(np.ascontiguousarray(outputs[1])).to(torch.int64).to(DEVICE)
-            mesh = Mesh(v=vertices, f=faces.to(torch.int64), device=DEVICE)
-            mesh.auto_normal()
+        vertices, faces = torch.from_numpy(outputs[0].astype(np.float32)).to(DEVICE), torch.from_numpy(np.ascontiguousarray(outputs[1])).to(torch.int64).to(DEVICE)
+        mesh = Mesh(v=vertices, f=faces.to(torch.int64), device=DEVICE)
+        mesh.auto_normal()
 
         return (mesh,)
 
@@ -4781,3 +4787,429 @@ class StableGen_StableX_Process_Image:
             
         except Exception as e:
             raise Exception(f"[StableGen_StableX_Process_Image] Image processing failed: {str(e)}")
+
+class Load_MVAdapter_IG2MV_Pipeline:
+    """Loader pipeline for MV-Adapter (Image to Multi-View)"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("DIFFUSERS_PIPE",)
+    RETURN_NAMES = ("mvadapter_pipe",)
+    FUNCTION = "load"
+
+    CKPT_MVADAPTER_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh", "MV-Adapter")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_model": (["stabilityai/stable-diffusion-xl-base-1.0"], 
+                             {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
+                "vae_model": (["madebyollin/sdxl-vae-fp16-fix", "None"], 
+                             {"default": "madebyollin/sdxl-vae-fp16-fix"}),
+                "adapter_path": (["huanngzh/mv-adapter"], {"default": "huanngzh/mv-adapter"}),
+                "scheduler": (["default", "ddpm", "lcm"], {"default": "default"}),
+                "num_views": ("INT", {"default": 6, "min": 1, "max": 16}),
+                "use_fp16": ("BOOLEAN", {"default": True}),
+                "use_mmgp": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "unet_model": ("STRING", {"default": ""}),
+                "lora_model": ("STRING", {"default": ""}),
+            }
+        }
+    @classmethod
+    def load(cls, base_model, vae_model, adapter_path, scheduler, num_views, 
+            use_fp16, use_mmgp, unet_model="", lora_model=""):
+        
+        
+        dtype = torch.float16 if use_fp16 else torch.float32
+        vae_model = None if vae_model == "None" else vae_model
+        unet_model = None if not unet_model else unet_model
+        lora_model = None if not lora_model else lora_model
+        
+        pipe = mvadapter_prepare_pipeline(
+            base_model=base_model,
+            vae_model=vae_model,
+            unet_model=unet_model,
+            lora_model=lora_model,
+            adapter_path=adapter_path,
+            scheduler=scheduler,
+            num_views=num_views,
+            device=DEVICE_STR,
+            dtype=dtype,
+            use_mmgp=use_mmgp,
+            adapter_local_path=cls.CKPT_MVADAPTER_PATH
+        )
+        
+        print("MV-Adapter pipeline loaded successfully")
+        return (pipe,)
+            
+
+class MVAdapter_Image_To_MultiView:
+    """Generate multi-view images from single image and 3D mesh"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("multiview_images",)
+    FUNCTION = "run"
+
+    @classmethod  
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mvadapter_pipe": ("DIFFUSERS_PIPE",),
+                "mesh_path": ("STRING", {"default": ""}),
+                "reference_image": ("IMAGE",),
+                "prompt": ("STRING", {"default": "high quality", "multiline": True}),
+                "negative_prompt": ("STRING", {"default": "watermark, ugly, deformed, noisy, blurry, low contrast", "multiline": True}),
+                "num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200}),
+                "guidance_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.1}),
+                "reference_conditioning_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
+                "height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
+                "width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
+                "seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
+                "remove_background": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
+            }
+        }
+
+    def run(self, mvadapter_pipe, mesh_path, reference_image, prompt, negative_prompt, 
+            num_inference_steps, guidance_scale, reference_conditioning_scale,
+            height, width, seed, remove_background, lora_scale=1.0):
+        
+        if isinstance(reference_image, torch.Tensor):
+            reference_images = torch_imgs_to_pils(reference_image)
+            reference_image = reference_images[0]
+        
+        if not mesh_path or not os.path.exists(mesh_path):
+            raise ValueError(f"Mesh path does not exist: {mesh_path}")
+        
+        num_views = 6 
+        images, pos_images, normal_images, processed_ref_image = mvadapter_run_pipeline(
+            pipe=mvadapter_pipe,
+            mesh_path=mesh_path,
+            num_views=num_views,
+            text=prompt,
+            image=reference_image,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            reference_conditioning_scale=reference_conditioning_scale,
+            negative_prompt=negative_prompt,
+            lora_scale=lora_scale,
+            device=DEVICE_STR,
+        )
+        
+        # grid_image = make_image_grid(images, rows=1)
+        
+        # grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
+        
+        # print(f"Generated multiview images: {grid_tensor.shape}")
+        # return (grid_tensor,)
+        # return images
+
+        return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
+        return (return_images,)
+
+
+class Load_MVAdapter_T2MV_Pipeline:
+    """Loader pipeline for MV-Adapter Text-to-Multi-View"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("DIFFUSERS_PIPE",)
+    RETURN_NAMES = ("mvadapter_t2mv_pipe",)
+    FUNCTION = "load"
+
+    CKPT_MVADAPTER_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh", "MV-Adapter")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_model": (["stabilityai/stable-diffusion-xl-base-1.0"], 
+                             {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
+                "vae_model": (["madebyollin/sdxl-vae-fp16-fix", "None"], 
+                             {"default": "madebyollin/sdxl-vae-fp16-fix"}),
+                "adapter_path": (["huanngzh/mv-adapter"], {"default": "huanngzh/mv-adapter"}),
+                "scheduler": (["default", "ddpm", "lcm"], {"default": "default"}),
+                "num_views": ("INT", {"default": 6, "min": 1, "max": 16}),
+                "use_fp16": ("BOOLEAN", {"default": True}),
+                "use_mmgp": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "unet_model": ("STRING", {"default": ""}),
+                "lora_model": ("STRING", {"default": ""}),
+            }
+        }
+
+    @classmethod
+    def load(cls, base_model, vae_model, adapter_path, scheduler, num_views, 
+             use_fp16, use_mmgp, unet_model="", lora_model=""):
+        
+        dtype = torch.float16 if use_fp16 else torch.float32
+        vae_model = None if vae_model == "None" else vae_model
+        unet_model = None if not unet_model else unet_model
+        lora_model = None if not lora_model else lora_model
+        
+        pipe = mvadapter_prepare_t2mv_pipeline(
+            base_model=base_model,
+            vae_model=vae_model,
+            unet_model=unet_model,
+            lora_model=lora_model,
+            adapter_path=adapter_path,
+            scheduler=scheduler,
+            num_views=num_views,
+            device=DEVICE_STR,
+            dtype=dtype,
+            use_mmgp=use_mmgp,
+            adapter_local_path=cls.CKPT_MVADAPTER_PATH
+        )
+        
+        print("MV-Adapter T2MV pipeline loaded successfully")
+        return (pipe,)
+
+class MVAdapter_Text_To_MultiView:
+    """Generate multi-view images from text prompt"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("multiview_images",)
+    FUNCTION = "run"
+
+    @classmethod  
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mvadapter_t2mv_pipe": ("DIFFUSERS_PIPE",),
+                "prompt": ("STRING", {"default": "a high quality 3D model", "multiline": True}),
+                "negative_prompt": ("STRING", {"default": "watermark, ugly, deformed, noisy, blurry, low contrast", "multiline": True}),
+                "num_views": ("INT", {"default": 6, "min": 1, "max": 16}),
+                "num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200}),
+                "guidance_scale": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 20.0, "step": 0.1}),
+                "height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
+                "width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
+                "seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
+            },
+            "optional": {
+                "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
+                "azimuth_angles": ("STRING", {"default": "0,45,90,180,270,315"}),
+            }
+        }
+
+    def run(self, mvadapter_t2mv_pipe, prompt, negative_prompt, num_views,
+            num_inference_steps, guidance_scale, height, width, seed, 
+            lora_scale=1.0, azimuth_angles="0,45,90,180,270,315"):
+        
+        # Parse azimuth angles
+        try:
+            azimuth_deg = [int(x.strip()) for x in azimuth_angles.split(",")]
+            # Ensure we have enough angles for the number of views
+            if len(azimuth_deg) < num_views:
+                # Repeat the pattern if we don't have enough angles
+                while len(azimuth_deg) < num_views:
+                    azimuth_deg.extend(azimuth_deg[:min(len(azimuth_deg), num_views - len(azimuth_deg))])
+            azimuth_deg = azimuth_deg[:num_views]  # Take only the first num_views
+        except:
+            # Fallback to default angles
+            azimuth_deg = [0, 45, 90, 180, 270, 315][:num_views]
+        
+        # Execute generation
+        images = mvadapter_run_t2mv_pipeline(
+            pipe=mvadapter_t2mv_pipe,
+            num_views=num_views,
+            text=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            negative_prompt=negative_prompt,
+            lora_scale=lora_scale,
+            device=DEVICE_STR,
+            azimuth_deg=azimuth_deg,
+        )
+        
+        # Create image grid  
+        # grid_image = make_image_grid(images, rows=1)
+        
+        # # Convert PIL image to torch tensor using shared utils
+        # grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
+        
+        # print(f"Generated T2MV multiview images: {grid_tensor.shape}")
+        # return (grid_tensor,)
+        return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
+        return (return_images,)
+            
+class Load_MVAdapter_Texture_Pipeline:
+    """Load texture projection pipeline for MVAdapter"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("TEXTURE_PIPE",)
+    RETURN_NAMES = ("texture_pipeline",)
+    FUNCTION = "load"
+
+    TEXTURE_CKPT_DIR = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh/MV-Adapter")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "upscaler_ckpt_name": ("STRING", {"default": "RealESRGAN_x2plus.pth"}),
+                "inpaint_ckpt_name": ("STRING", {"default": "big-lama.pt"}),
+                "use_mmgp": ("BOOLEAN", {"default": False}),
+                "auto_download": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    @classmethod
+    def load(cls, upscaler_ckpt_name, inpaint_ckpt_name, use_mmgp, auto_download):
+        
+        # Form full paths from model names
+        upscaler_ckpt_path = os.path.join(cls.TEXTURE_CKPT_DIR, upscaler_ckpt_name) if upscaler_ckpt_name.strip() else None
+        inpaint_ckpt_path = os.path.join(cls.TEXTURE_CKPT_DIR, inpaint_ckpt_name) if inpaint_ckpt_name.strip() else None
+        
+        # Auto-download checkpoints if enabled
+        if auto_download:
+            download_texture_checkpoints(cls.TEXTURE_CKPT_DIR, upscaler_ckpt_path, inpaint_ckpt_path)
+        
+        texture_pipe = mvadapter_prepare_texture_pipeline(
+            upscaler_ckpt_path=upscaler_ckpt_path,
+            inpaint_ckpt_path=inpaint_ckpt_path,
+            device=DEVICE_STR,
+            use_mmgp=use_mmgp,
+        )
+        
+        print("Texture pipeline loaded successfully")
+        return (texture_pipe,)
+    
+
+class MVAdapter_Texture_Projection:
+    """Project grid image onto 3D mesh using pre-loaded texture pipeline"""
+    CATEGORY = "Comfy3D/Algorithm"
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("shaded_model_path", "pbr_model_path") 
+    FUNCTION = "run"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "texture_pipeline": ("TEXTURE_PIPE",),
+                "grid_image": ("IMAGE",),
+                "mesh_path": ("STRING", {"default": ""}),
+                "save_dir": ("STRING", {"default": "./output"}),
+                "save_name": ("STRING", {"default": "textured_model"}),
+                "uv_size": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 256}),
+                "view_upscale": ("BOOLEAN", {"default": True}),
+                "inpaint_mode": (["none", "uv", "view"], {"default": "view"}),
+                "uv_unwarp": ("BOOLEAN", {"default": True}),
+                "preprocess_mesh": ("BOOLEAN", {"default": False}),
+                "move_to_center": ("BOOLEAN", {"default": False}),
+                "front_x": ("BOOLEAN", {"default": True}),
+                "create_pbr_model": ("BOOLEAN", {"default": True}),
+                "apply_dilate": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "camera_azimuth_deg": ("STRING", {"default": "0,90,180,270,180,180"}),
+                "camera_elevation_deg": ("STRING", {"default": "0,0,0,0,89.99,-89.99"}),
+                "camera_distance": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
+                "camera_ortho_scale": ("FLOAT", {"default": 1.1, "min": 0.1, "max": 5.0, "step": 0.1}),
+                "debug_mode": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    def run(self, texture_pipeline, grid_image, mesh_path, save_dir, save_name, 
+            uv_size, view_upscale, inpaint_mode, uv_unwarp, preprocess_mesh, 
+            move_to_center, front_x, create_pbr_model, apply_dilate,
+            camera_azimuth_deg="0,90,180,270,180,180",
+            camera_elevation_deg="0,0,0,0,89.99,-89.99", 
+            camera_distance=1.0, camera_ortho_scale=1.1, debug_mode=False):
+        
+        if isinstance(grid_image, torch.Tensor):
+            pil_grid = torch_imgs_to_pils(grid_image)[0]  # Get first (and only) image
+        else:
+            raise ValueError("grid_image must be torch.Tensor")
+        
+        if not mesh_path or not os.path.exists(mesh_path):
+            raise ValueError(f"Mesh file does not exist: {mesh_path}")
+        
+        # Convert to absolute paths
+        save_dir = os.path.abspath(save_dir)
+        os.makedirs(save_dir, exist_ok=True)
+        
+        # Parse camera angles with fallback to defaults
+        try:
+            azimuth_deg = [float(x.strip()) for x in camera_azimuth_deg.split(",")]
+            elevation_deg = [float(x.strip()) for x in camera_elevation_deg.split(",")]
+        except:
+            azimuth_deg = [0, 90, 180, 270, 180, 180]
+            elevation_deg = [0, 0, 0, 0, 89.99, -89.99]
+        
+        # Adjust azimuth angles (subtract 90 as in original code)
+        azimuth_deg_corrected = [x - 90 for x in azimuth_deg]
+        
+        # Save grid image temporarily
+        temp_grid_path = os.path.join(save_dir, f"{save_name}_temp_grid.png")
+        pil_grid.save(temp_grid_path)
+        
+        try:
+            from .Gen_3D_Modules.MV_Adapter.mvadapter.pipelines.pipeline_texture import ModProcessConfig
+            
+            rgb_process_config = ModProcessConfig(
+                view_upscale=view_upscale,
+                inpaint_mode=inpaint_mode
+            )
+            
+            base_color_process_config = ModProcessConfig(
+                view_upscale=view_upscale,
+                inpaint_mode=inpaint_mode
+            )
+            
+            # Prepare texture pipeline arguments
+            texture_args = {
+                "mesh_path": mesh_path,
+                "save_dir": save_dir,
+                "save_name": save_name,
+                "move_to_center": move_to_center,
+                "front_x": front_x,
+                "uv_unwarp": uv_unwarp,
+                "preprocess_mesh": preprocess_mesh,
+                "uv_size": uv_size,
+                "rgb_path": temp_grid_path,
+                "rgb_process_config": rgb_process_config,
+                "camera_elevation_deg": elevation_deg,
+                "camera_azimuth_deg": azimuth_deg_corrected,
+                "camera_distance": camera_distance,
+                "camera_ortho_scale": camera_ortho_scale,
+                "debug_mode": debug_mode,
+                "apply_dilate": apply_dilate,
+            }
+            
+            # Add PBR settings if enabled
+            if create_pbr_model:
+                texture_args.update({
+                    "base_color_path": temp_grid_path,
+                    "base_color_process_config": base_color_process_config,
+                })
+            
+            output = texture_pipeline(**texture_args)
+            
+            # Cleanup temp file
+            if os.path.exists(temp_grid_path):
+                os.remove(temp_grid_path)
+            
+            # Convert to absolute paths
+            shaded_path = os.path.abspath(output.shaded_model_save_path) if output.shaded_model_save_path else ""
+            pbr_path = os.path.abspath(output.pbr_model_save_path) if output.pbr_model_save_path else ""
+            
+            print(f"Texturing from grid completed:")
+            print(f"  Shaded model: {shaded_path}")
+            print(f"  PBR model: {pbr_path}")
+                
+            return (shaded_path, pbr_path)
+            
+        except Exception as e:
+            # Cleanup temp file on error
+            if os.path.exists(temp_grid_path):
+                os.remove(temp_grid_path)
+            raise e
+            

pyproject.toml

@@ -1,5 +1,5 @@
 [project]
-name = "comfyui-3d-pack"
+name = "ComfyUI-3D-Pack"
 description = "Make ComfyUI generates 3D assets as good & convenient as it generates image/video!"
 version = "0.1.4"
 license = { file = "LICENSE" }