content: from dataclasses import dataclass, field
  
  import numpy as np
  import json
  import copy
  import torch
  import torch.nn.functional as F
  from skimage import measure
  from einops import repeat
  from tqdm import tqdm
  from PIL import Image
  
  from diffusers import (
      DDPMScheduler,
      DDIMScheduler,
      UniPCMultistepScheduler,
      KarrasVeScheduler,
      DPMSolverMultistepScheduler
  )
  
  import craftsman
  from craftsman.systems.base import BaseSystem
  from craftsman.utils.ops import generate_dense_grid_points
  from craftsman.utils.misc import get_rank
  from craftsman.utils.typing import *
  
  import comfy.utils
  
  def compute_snr(noise_scheduler, timesteps):
      """
      Computes SNR as per
      https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
      """
      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 https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
      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 ddim_sample(ddim_scheduler: DDIMScheduler,
                  diffusion_model: torch.nn.Module,
                  shape: Union[List[int], Tuple[int]],
                  cond: torch.FloatTensor,
                  steps: int,
                  eta: float = 0.0,
                  guidance_scale: float = 3.0,
                  do_classifier_free_guidance: bool = True,
                  generator: Optional[torch.Generator] = None,
                  device: torch.device = "cuda:0",
                  disable_prog: bool = True):
  
      assert steps > 0, f"{steps} must > 0."
  
      # init latents
      bsz = cond.shape[0]
      if do_classifier_free_guidance:
          bsz = bsz // 2
  
      latents = torch.randn(
          (bsz, *shape),
          generator=generator,
          device=cond.device,
          dtype=cond.dtype,
      )
      # scale the initial noise by the standard deviation required by the scheduler
      latents = latents * ddim_scheduler.init_noise_sigma
      # set timesteps
      ddim_scheduler.set_timesteps(steps)
      timesteps = ddim_scheduler.timesteps.to(device)
      # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
      # eta (η) is only used with the DDIMScheduler, and between [0, 1]
      extra_step_kwargs = {
          # "eta": eta,
          "generator": generator
      }
  
      # reverse
      comfy_pbar = comfy.utils.ProgressBar(steps)
      for i, t in enumerate(tqdm(timesteps, disable=disable_prog, desc="DDIM Sampling:", leave=False)):
          # expand the latents if we are doing classifier free guidance
          latent_model_input = (
              torch.cat([latents] * 2)
              if do_classifier_free_guidance
              else latents
          )
          # predict the noise residual
          timestep_tensor = torch.tensor([t], dtype=torch.long, device=device)
          timestep_tensor = timestep_tensor.expand(latent_model_input.shape[0])
          noise_pred = diffusion_model.forward(latent_model_input, timestep_tensor, cond)
  
          # perform guidance
          if do_classifier_free_guidance:
              noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
              noise_pred = noise_pred_uncond + guidance_scale * (
                      noise_pred_text - noise_pred_uncond
              )
  
          # compute the previous noisy sample x_t -> x_t-1
          latents = ddim_scheduler.step(
              noise_pred, t, latents, **extra_step_kwargs
          ).prev_sample
          
          comfy_pbar.update_absolute(i + 1)
  
          yield latents, t
  
  
  @craftsman.register("shape-diffusion-system")
  class ShapeDiffusionSystem(BaseSystem):
      @dataclass
      class Config(BaseSystem.Config):
          val_samples_json: str = None
          z_scale_factor: float = 1.0
          guidance_scale: float = 7.5
          num_inference_steps: int = 50
          eta: float = 0.0
          snr_gamma: float = 5.0
  
          # shape vae model
          shape_model_type: str = None
          shape_model: dict = field(default_factory=dict)
  
          # condition model
          condition_model_type: str = None
          condition_model: dict = field(default_factory=dict)
  
          # diffusion model
          denoiser_model_type: str = None
          denoiser_model: dict = field(default_factory=dict)
  
          # noise scheduler
          noise_scheduler_type: str = None
          noise_scheduler: dict = field(default_factory=dict)
  
          # denoise scheduler
          denoise_scheduler_type: str = None
          denoise_scheduler: dict = field(default_factory=dict)
  
      cfg: Config
  
      def configure(self):
          super().configure()
  
          self.shape_model = craftsman.find(self.cfg.shape_model_type)(self.cfg.shape_model)
          self.shape_model.eval()
          self.shape_model.requires_grad_(False)
  
          self.condition = craftsman.find(self.cfg.condition_model_type)(self.cfg.condition_model)
  
          self.denoiser_model = craftsman.find(self.cfg.denoiser_model_type)(self.cfg.denoiser_model)
  
          self.noise_scheduler = craftsman.find(self.cfg.noise_scheduler_type)(**self.cfg.noise_scheduler)
          self.denoise_scheduler = craftsman.find(self.cfg.denoise_scheduler_type)(**self.cfg.denoise_scheduler)
  
          self.z_scale_factor = self.cfg.z_scale_factor
  
      def forward(self, batch: Dict[str, Any]) -> Dict[str, Any]:
          # encode shape latents
          shape_embeds, kl_embed, posterior = self.shape_model.encode(
              batch["surface"][..., :3 + self.cfg.shape_model.point_feats], 
              sample_posterior=True
          )
          latents = kl_embed * self.z_scale_factor
  
          cond_latents = self.condition(batch)
          cond_latents = cond_latents.to(latents).view(latents.shape[0], -1, cond_latents.shape[-1])
  
          # Sample noise that we"ll add to the latents
          # [batch_size, n_token, latent_dim]
          noise = torch.randn_like(latents).to(latents)
          bs = latents.shape[0]
          # Sample a random timestep for each motion
          timesteps = torch.randint(
              0,
              self.noise_scheduler.config.num_train_timesteps,
              (bs,),
              device=latents.device,
          )
          # import pdb; pdb.set_trace()
  
          timesteps = timesteps.long()
          # Add noise to the latents according to the noise magnitude at each timestep
          # x_t
          noisy_z = self.noise_scheduler.add_noise(latents, noise, timesteps)
          
          # diffusion model forward
          noise_pred = self.denoiser_model(noisy_z, timesteps, cond_latents)
  
          # compute loss
          if self.noise_scheduler.config.prediction_type == "epsilon":
              target = noise 
          elif self.noise_scheduler.config.prediction_type == "v_prediction":
              target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
          else:
              raise ValueError(f"Prediction Type: {self.noise_scheduler.prediction_type} not supported.")
          if self.cfg.snr_gamma == 0:
              if self.cfg.loss.loss_type == "l1":
                  loss = F.l1_loss(noise_pred, target, reduction="mean")
              elif self.cfg.loss.loss_type in ["mse", "l2"]:
                  loss = F.mse_loss(noise_pred, target, reduction="mean")
              else:
                  raise ValueError(f"Loss Type: {self.cfg.loss.loss_type} not supported.")
          else:
              # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
              # Since we predict the noise instead of x_0, the original formulation is slightly changed.
              # This is discussed in Section 4.2 of the same paper.
              snr = compute_snr(self.noise_scheduler, timesteps)
              mse_loss_weights = torch.stack([snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1).min(
                  dim=1
              )[0]
              if self.noise_scheduler.config.prediction_type == "epsilon":
                  mse_loss_weights = mse_loss_weights / snr
              elif noise_scheduler.config.prediction_type == "v_prediction":
                  mse_loss_weights = mse_loss_weights / (snr + 1)
              
              if self.cfg.loss.loss_type == "l1":
                  loss = F.l1_loss(noise_pred, target, reduction="none")
              elif self.cfg.loss.loss_type in ["mse", "l2"]:
                  loss = F.mse_loss(noise_pred, target, reduction="none")
              else:
                  raise ValueError(f"Loss Type: {self.cfg.loss.loss_type} not supported.")
              loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
              loss = loss.mean()
  
          return {
              "loss_diffusion": loss,
              "latents": latents,
              "x_0": noisy_z,
              "noise": noise,
              "noise_pred": noise_pred,
              "timesteps": timesteps,
              }
  
      def training_step(self, batch, batch_idx):
          out = self(batch)
  
          loss = 0.
          for name, value in out.items():
              if name.startswith("loss_"):
                  self.log(f"train/{name}", value)
                  loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")])
  
          for name, value in self.cfg.loss.items():
              if name.startswith("lambda_"):
                  self.log(f"train_params/{name}", self.C(value))
  
          return {"loss": loss}
  
      @torch.no_grad()
      def validation_step(self, batch, batch_idx):
          self.eval()
  
          if get_rank() == 0:
              sample_inputs = json.loads(open(self.cfg.val_samples_json).read()) # condition
              sample_inputs_ = copy.deepcopy(sample_inputs)
              sample_outputs = self.sample(sample_inputs) # list
              for i, sample_output in enumerate(sample_outputs):
                  mesh_v_f, has_surface = self.shape_model.extract_geometry(sample_output, octree_depth=7)
                  for j in range(len(mesh_v_f)):
                      if "text" in sample_inputs_ and "image" in sample_inputs_:
                          name = sample_inputs_["image"][j].split("/")[-1].replace(".png", "")
                      elif "text" in sample_inputs_ and "mvimage" in sample_inputs_:
                          name = sample_inputs_["mvimages"][j][0].split("/")[-2].replace(".png", "")
                      elif "text" in sample_inputs_:
                          name = sample_inputs_["text"][j].replace(" ", "_")
                      elif "image" in sample_inputs_:
                          name = sample_inputs_["image"][j].split("/")[-1].replace(".png", "")
                      elif "mvimages" in sample_inputs_:
                          name = sample_inputs_["mvimages"][j][0].split("/")[-2].replace(".png", "")
                      self.save_mesh(
                          f"it{self.true_global_step}/{name}_{i}.obj",
                          mesh_v_f[j][0], mesh_v_f[j][1]
                      )
  
          out = self(batch)
          if self.global_step == 0:
              latents = self.shape_model.decode(out["latents"])
              mesh_v_f, has_surface = self.shape_model.extract_geometry(latents)
              self.save_mesh(
                  f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.obj",
                  mesh_v_f[0][0], mesh_v_f[0][1]
              )
              # exit()
          torch.cuda.empty_cache()
  
          return {"val/loss": out["loss_diffusion"]}
  
      @torch.no_grad()
      def sample(self,
                 sample_inputs: Dict[str, Union[torch.FloatTensor, List[str]]],
                 sample_times: int = 1,
                 steps: Optional[int] = None,
                 guidance_scale: Optional[float] = None,
                 eta: float = 0.0,
                 return_intermediates: bool = False, 
                 camera_embeds: Optional[torch.Tensor] = None,
                 seed: Optional[int] = None,
                 **kwargs):
  
          if steps is None:
              steps = self.cfg.num_inference_steps
          if guidance_scale is None:
              guidance_scale = self.cfg.guidance_scale
          do_classifier_free_guidance = guidance_scale > 0
  
          # conditional encode
          if "image" in sample_inputs:
              sample_inputs["image"] = [Image.open(img) for img in sample_inputs["image"]]
              cond = self.condition.encode_image(sample_inputs["image"])
              if do_classifier_free_guidance:
                  un_cond = self.condition.empty_image_embeds.repeat(len(sample_inputs["image"]), 1, 1).to(cond)
                  cond = torch.cat([un_cond, cond], dim=0)
          elif "mvimages" in sample_inputs: # by default 4 views
              bs = len(sample_inputs["mvimages"])
              cond = []
              for image in sample_inputs["mvimages"]:
                  if isinstance(image, list) and isinstance(image[0], str):
                      sample_inputs["image"] = [Image.open(img) for img in image] # List[PIL]
                  else:
                      sample_inputs["image"] = image
                  cond += [self.condition.encode_image(sample_inputs["image"])]
              cond = torch.stack(cond, dim=0)# tensor  shape 为[len(sample_inputs["mvimages"], 4*(num_latents+1), context_dim]
              if do_classifier_free_guidance:
                  un_cond = self.condition.empty_image_embeds.unsqueeze(0).repeat(len(sample_inputs["mvimages"]), cond.shape[1] // self.condition.cfg.n_views, 1, 1).to(cond) # shape 为[len(sample_inputs["mvimages"], 4*(num_latents+1), context_dim]
                  cond = torch.cat([un_cond, cond], dim=0).view(bs * 2, -1, cond[0].shape[-1]) 
          else:
              raise NotImplementedError("Only text, image or mvimages condition is supported.")
  
          outputs = []
          latents = None
          
          if seed != None:
              generator = torch.Generator(device="cuda").manual_seed(seed)
          else:
              generator = None
              
          if not return_intermediates:
              for _ in range(sample_times):
                  sample_loop = ddim_sample(
                      self.denoise_scheduler,
                      self.denoiser_model.eval(),
                      shape=self.shape_model.latent_shape,
                      cond=cond,
                      steps=steps,
                      guidance_scale=guidance_scale,
                      do_classifier_free_guidance=do_classifier_free_guidance,
                      device=self.device,
                      eta=eta,
                      disable_prog=False,
                      generator= generator
                  )
                  for sample, t in sample_loop:
                      latents = sample
                  outputs.append(self.shape_model.decode(latents / self.z_scale_factor, **kwargs))
          else:
              sample_loop = ddim_sample(
                  self.denoise_scheduler,
                  self.denoiser_model.eval(),
                  shape=self.shape_model.latent_shape,
                  cond=cond,
                  steps=steps,
                  guidance_scale=guidance_scale,
                  do_classifier_free_guidance=do_classifier_free_guidance,
                  device=self.device,
                  eta=eta,
                  disable_prog=False,
                  generator= generator
              )
  
              iter_size = steps // sample_times
              i = 0
              for sample, t in sample_loop:
                  latents = sample
                  if i % iter_size == 0 or i == steps - 1:
                      outputs.append(self.shape_model.decode(latents / self.z_scale_factor, **kwargs))
                  i += 1
  
          return outputs
      
  
      def on_validation_epoch_end(self):
          pass
  