i

.gitignore

@@ -1 +1,2 @@
-gradio_cached_dir
+gradio_cached_dir
+jiangxin

ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml

@@ -0,0 +1,143 @@
+name: michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k
+description: ''
+tag: michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue
+seed: 0
+use_timestamp: true
+timestamp: ''
+exp_root_dir: outputs
+exp_dir: outputs/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k
+trial_name: michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue
+trial_dir: outputs/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k/michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue
+n_gpus: 8
+resume: ./ckpts/3DNativeGeneration/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k.ckpt
+data_type: objaverse-datamodule
+data:
+  root_dir: data/objaverse_clean/cap3d_high_quality_170k_images
+  data_type: occupancy
+  n_samples: 4096
+  noise_sigma: 0.0
+  load_supervision: false
+  supervision_type: occupancy
+  n_supervision: 10000
+  load_image: true
+  image_data_path: data/objaverse_clean/raw_data/images/cap3d_high_quality_170k
+  image_type: mvrgb
+  idx:
+  - 0
+  - 4
+  - 8
+  - 12
+  - 16
+  n_views: 4
+  load_caption: false
+  rotate_points: false
+  batch_size: 32
+  num_workers: 16
+system_type: shape-diffusion-system
+system:
+  val_samples_json: val_data/mv_images/val_samples_rgb_mvimage.json
+  z_scale_factor: 1.0
+  guidance_scale: 7.5
+  num_inference_steps: 50
+  eta: 0.0
+  shape_model_type: michelangelo-aligned-autoencoder
+  shape_model:
+    num_latents: 256
+    embed_dim: 64
+    point_feats: 3
+    out_dim: 1
+    num_freqs: 8
+    include_pi: false
+    heads: 12
+    width: 768
+    num_encoder_layers: 8
+    num_decoder_layers: 16
+    use_ln_post: true
+    init_scale: 0.25
+    qkv_bias: false
+    use_flash: true
+    use_checkpoint: true
+  condition_model_type: clip-embedder
+  condition_model:
+    pretrained_model_name_or_path: openai/clip-vit-large-patch14
+    encode_camera: true
+    camera_embeds_dim: 32
+    n_views: 4
+    empty_embeds_ratio: 0.1
+    normalize_embeds: false
+    zero_uncond_embeds: true
+  denoiser_model_type: simple-denoiser
+  denoiser_model:
+    input_channels: 64
+    output_channels: 64
+    n_ctx: 256
+    width: 768
+    layers: 6
+    heads: 12
+    context_dim: 1024
+    init_scale: 1.0
+    skip_ln: true
+    use_checkpoint: true
+  noise_scheduler_type: diffusers.schedulers.DDPMScheduler
+  noise_scheduler:
+    num_train_timesteps: 1000
+    beta_start: 0.00085
+    beta_end: 0.012
+    beta_schedule: scaled_linear
+    variance_type: fixed_small
+    clip_sample: false
+  denoise_scheduler_type: diffusers.schedulers.DDIMScheduler
+  denoise_scheduler:
+    num_train_timesteps: 1000
+    beta_start: 0.00085
+    beta_end: 0.012
+    beta_schedule: scaled_linear
+    clip_sample: false
+    set_alpha_to_one: false
+    steps_offset: 1
+  loggers:
+    wandb:
+      enable: false
+      project: JiangXin
+      name: text-to-shape-diffusion+michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k+michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue
+  loss:
+    loss_type: mse
+    lambda_diffusion: 1.0
+  optimizer:
+    name: AdamW
+    args:
+      lr: 5.0e-05
+      betas:
+      - 0.9
+      - 0.99
+      eps: 1.0e-06
+  scheduler:
+    name: SequentialLR
+    interval: step
+    schedulers:
+    - name: LinearLR
+      interval: step
+      args:
+        start_factor: 1.0e-06
+        end_factor: 1.0
+        total_iters: 5000
+    - name: CosineAnnealingLR
+      interval: step
+      args:
+        T_max: 5000
+        eta_min: 0.0
+    milestones:
+    - 5000
+trainer:
+  num_nodes: 2
+  max_epochs: 100000
+  log_every_n_steps: 5
+  num_sanity_val_steps: 1
+  check_val_every_n_epoch: 3
+  enable_progress_bar: true
+  precision: 16-mixed
+  strategy: ddp_find_unused_parameters_true
+checkpoint:
+  save_last: true
+  save_top_k: -1
+  every_n_train_steps: 5000

ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41248dba953cad356c491e7584b4171920f2ad95af10b0f78225eda867dbb7c4
+size 3722911570

craftsman/models/autoencoders/michelangelo_autoencoder.py

@@ -324,3 +324,81 @@ class MichelangeloAutoencoder(AutoEncoder):
         logits = self.decoder(queries, latents).squeeze(-1)
 
         return logits
+
+
+
+
+@craftsman.register("michelangelo-aligned-autoencoder")
+class MichelangeloAlignedAutoencoder(MichelangeloAutoencoder):
+    r"""
+    A VAE model for encoding shapes into latents and decoding latent representations into shapes.
+    """
+    @dataclass
+    class Config(MichelangeloAutoencoder.Config):
+        clip_model_version: Optional[str] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        if self.cfg.clip_model_version is not None:
+            self.clip_model: CLIPModel = CLIPModel.from_pretrained(self.cfg.clip_model_version)
+            self.projection = nn.Parameter(torch.empty(self.cfg.width, self.clip_model.projection_dim))
+            self.logit_scale = torch.exp(self.clip_model.logit_scale.data)
+            nn.init.normal_(self.projection, std=self.clip_model.projection_dim ** -0.5)
+        else:
+            self.projection = nn.Parameter(torch.empty(self.cfg.width, 768))
+            nn.init.normal_(self.projection, std=768 ** -0.5)
+
+        self.cfg.num_latents = self.cfg.num_latents + 1
+
+        super().configure()
+
+    def encode(self,
+               surface: torch.FloatTensor,
+               sample_posterior: bool = True):
+        """
+        Args:
+            surface (torch.FloatTensor): [B, N, 3+C]
+            sample_posterior (bool):
+
+        Returns:
+            latents (torch.FloatTensor)
+            posterior (DiagonalGaussianDistribution or None):
+        """
+        assert surface.shape[-1] == 3 + self.cfg.point_feats, f"\
+            Expected {3 + self.cfg.point_feats} channels, got {surface.shape[-1]}"
+        
+        pc, feats = surface[..., :3], surface[..., 3:] # B, n_samples, 3    
+        shape_latents = self.encoder(pc, feats) # B, num_latents, width
+        shape_embeds = shape_latents[:, 0]  # B, width
+        shape_latents = shape_latents[:, 1:] # B, num_latents-1, width
+        kl_embed, posterior = self.encode_kl_embed(shape_latents, sample_posterior)  # B, num_latents, embed_dim
+
+        shape_embeds = shape_embeds @ self.projection
+        return shape_embeds, kl_embed, posterior
+    
+    def forward(self,
+                surface: torch.FloatTensor,
+                queries: torch.FloatTensor,
+                sample_posterior: bool = True):
+        """
+        Args:
+            surface (torch.FloatTensor): [B, N, 3+C]
+            queries (torch.FloatTensor): [B, P, 3]
+            sample_posterior (bool):
+
+        Returns:
+            shape_embeds (torch.FloatTensor): [B, width]
+            latents (torch.FloatTensor): [B, num_latents, embed_dim]
+            posterior (DiagonalGaussianDistribution or None).
+            logits (torch.FloatTensor): [B, P]
+        """
+
+        shape_embeds, kl_embed, posterior = self.encode(surface, sample_posterior=sample_posterior)
+
+        latents = self.decode(kl_embed) # [B, num_latents - 1, width]
+
+        logits = self.query(queries, latents) # [B,]
+
+        return shape_embeds, latents, posterior, logits
+    

gradio_app.py

@@ -170,8 +170,10 @@ if __name__=="__main__":
     # mvimg_model_config_list = ["CRM", "ImageDream", "Wonder3D"]
     
     # for 3D latent set diffusion
-    ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt", repo_type="model")
-    config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml", repo_type="model")
+    ckpt_path = "./ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt"
+    config_path = "./ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml"
+    # ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt", repo_type="model")
+    # config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml", repo_type="model")
     scheluder_dict = OrderedDict({ 
         "DDIMScheduler": 'diffusers.schedulers.DDIMScheduler',
         # "DPMSolverMultistepScheduler": 'diffusers.schedulers.DPMSolverMultistepScheduler', # not support yet