content: import os
  from contextlib import contextmanager
  from typing import List, Tuple, Optional, Union
  
  import torch
  import torch.nn as nn
  from torch.optim import lr_scheduler
  import pytorch_lightning as pl
  from pytorch_lightning.utilities import rank_zero_info
  from pytorch_lightning.utilities import rank_zero_only
  
  from ...utils.ema import LitEma
  from ...utils.misc import instantiate_from_config, instantiate_non_trainable_model
  
  
  
  class Diffuser(pl.LightningModule):
      def __init__(
          self,
          *,
          first_stage_config,
          cond_stage_config,
          denoiser_cfg,
          scheduler_cfg,
          optimizer_cfg,
          pipeline_cfg=None,
          image_processor_cfg=None,
          lora_config=None,
          ema_config=None,
          first_stage_key: str = "surface",
          cond_stage_key: str = "image",
          scale_by_std: bool = False,
          z_scale_factor: float = 1.0,
          ckpt_path: Optional[str] = None,
          ignore_keys: Union[Tuple[str], List[str]] = (),
          torch_compile: bool = False,
      ):
          super().__init__()
          self.first_stage_key = first_stage_key
          self.cond_stage_key = cond_stage_key
  
          # ========= init optimizer config ========= #
          self.optimizer_cfg = optimizer_cfg
  
          # ========= init diffusion scheduler ========= #
          self.scheduler_cfg = scheduler_cfg
          self.sampler = None
          if 'transport' in scheduler_cfg:
              self.transport = instantiate_from_config(scheduler_cfg.transport)
              self.sampler = instantiate_from_config(scheduler_cfg.sampler, transport=self.transport)
              self.sample_fn = self.sampler.sample_ode(**scheduler_cfg.sampler.ode_params)
  
          # ========= init the model ========= #
          self.denoiser_cfg = denoiser_cfg
          self.model = instantiate_from_config(denoiser_cfg, device=None, dtype=None)
          self.cond_stage_model = instantiate_from_config(cond_stage_config)
  
          self.ckpt_path = ckpt_path
          if ckpt_path is not None:
              self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
  
          # ========= config lora model ========= #
          if lora_config is not None:
              from peft import LoraConfig, get_peft_model
              loraconfig = LoraConfig(
                  r=lora_config.rank,
                  lora_alpha=lora_config.rank,
                  target_modules=lora_config.get('target_modules')
              )
              self.model = get_peft_model(self.model, loraconfig)
  
          # ========= config ema model ========= #
          self.ema_config = ema_config
          if self.ema_config is not None:
              if self.ema_config.ema_model == 'DSEma':
                  # from michelangelo.models.modules.ema_deepspeed import DSEma
                  from ..utils.ema_deepspeed import DSEma
                  self.model_ema = DSEma(self.model, decay=self.ema_config.ema_decay)
              else:
                  self.model_ema = LitEma(self.model, decay=self.ema_config.ema_decay)
              #do not initilize EMA weight from ckpt path, since I need to change moe layers
              if ckpt_path is not None:
                  self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
              print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
  
          # ========= init vae at last to prevent it is overridden by loaded ckpt ========= #
          self.first_stage_model = instantiate_non_trainable_model(first_stage_config)
  
          self.scale_by_std = scale_by_std
          if scale_by_std:
              self.register_buffer("z_scale_factor", torch.tensor(z_scale_factor))
          else:
              self.z_scale_factor = z_scale_factor
  
          # ========= init pipeline for inference ========= #
          self.image_processor_cfg = image_processor_cfg
          self.image_processor = None
          if self.image_processor_cfg is not None:
              self.image_processor = instantiate_from_config(self.image_processor_cfg)
          self.pipeline_cfg = pipeline_cfg
          from ...schedulers import FlowMatchEulerDiscreteScheduler
          scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
          self.pipeline = instantiate_from_config(
              pipeline_cfg,
              vae=self.first_stage_model,
              model=self.model,
              scheduler=scheduler, # self.sampler,
              conditioner=self.cond_stage_model,
              image_processor=self.image_processor,
          )
  
          # ========= torch compile to accelerate ========= #
          self.torch_compile = torch_compile
          if self.torch_compile:
              torch.nn.Module.compile(self.model)
              torch.nn.Module.compile(self.first_stage_model)
              torch.nn.Module.compile(self.cond_stage_model)
              print(f'*' * 100)
              print(f'Compile model for acceleration')
              print(f'*' * 100)
  
      @contextmanager
      def ema_scope(self, context=None):
          if self.ema_config is not None and self.ema_config.get('ema_inference', False):
              self.model_ema.store(self.model)
              self.model_ema.copy_to(self.model)
              if context is not None:
                  print(f"{context}: Switched to EMA weights")
          try:
              yield None
          finally:
              if self.ema_config is not None and self.ema_config.get('ema_inference', False):
                  self.model_ema.restore(self.model)
                  if context is not None:
                      print(f"{context}: Restored training weights")
  
      def init_from_ckpt(self, path, ignore_keys=()):
          ckpt = torch.load(path, map_location="cpu")
          if 'state_dict' not in ckpt:
              # deepspeed ckpt
              state_dict = {}
              for k in ckpt.keys():
                  new_k = k.replace('_forward_module.', '')
                  state_dict[new_k] = ckpt[k]
          else:
              state_dict = ckpt["state_dict"]
  
          keys = list(state_dict.keys())
          for k in keys:
              for ik in ignore_keys:
                  if ik in k:
                      print("Deleting key {} from state_dict.".format(k))
                      del state_dict[k]
  
          missing, unexpected = self.load_state_dict(state_dict, strict=False)
          print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
          if len(missing) > 0:
              print(f"Missing Keys: {missing}")
              print(f"Unexpected Keys: {unexpected}")
  
      def on_load_checkpoint(self, checkpoint):
          """
          The pt_model is trained separately, so we already have access to its
          checkpoint and load it separately with `self.set_pt_model`.
  
          However, the PL Trainer is strict about
          checkpoint loading (not configurable), so it expects the loaded state_dict
          to match exactly the keys in the model state_dict.
  
          So, when loading the checkpoint, before matching keys, we add all pt_model keys
          from self.state_dict() to the checkpoint state dict, so that they match
          """
          for key in self.state_dict().keys():
              if key.startswith("model_ema") and key not in checkpoint["state_dict"]:
                  checkpoint["state_dict"][key] = self.state_dict()[key]
  
      def configure_optimizers(self) -> Tuple[List, List]:
          lr = self.learning_rate
  
          params_list = []
          trainable_parameters = list(self.model.parameters())
          params_list.append({'params': trainable_parameters, 'lr': lr})
  
          no_decay = ['bias', 'norm.weight', 'norm.bias', 'norm1.weight', 'norm1.bias', 'norm2.weight', 'norm2.bias']
  
  
          if self.optimizer_cfg.get('train_image_encoder', False):
              image_encoder_parameters = list(self.cond_stage_model.named_parameters())
              image_encoder_parameters_decay = [param for name, param in image_encoder_parameters if
                                                not any((no_decay_name in name) for no_decay_name in no_decay)]
              image_encoder_parameters_nodecay = [param for name, param in image_encoder_parameters if
                                                  any((no_decay_name in name) for no_decay_name in no_decay)]
              # filter trainable params
              image_encoder_parameters_decay = [param for param in image_encoder_parameters_decay if
                                                param.requires_grad]
              image_encoder_parameters_nodecay = [param for param in image_encoder_parameters_nodecay if
                                                  param.requires_grad]
  
              print(f"Image Encoder Params: {len(image_encoder_parameters_decay)} decay, ")
              print(f"Image Encoder Params: {len(image_encoder_parameters_nodecay)} nodecay, ")
  
              image_encoder_lr = self.optimizer_cfg['image_encoder_lr']
              image_encoder_lr_multiply = self.optimizer_cfg.get('image_encoder_lr_multiply', 1.0)
              image_encoder_lr = image_encoder_lr if image_encoder_lr is not None else lr * image_encoder_lr_multiply
              params_list.append(
                  {'params': image_encoder_parameters_decay, 'lr': image_encoder_lr,
                   'weight_decay': 0.05})
              params_list.append(
                  {'params': image_encoder_parameters_nodecay, 'lr': image_encoder_lr,
                   'weight_decay': 0.})
  
          optimizer = instantiate_from_config(self.optimizer_cfg.optimizer, params=params_list, lr=lr)
          if hasattr(self.optimizer_cfg, 'scheduler'):
              scheduler_func = instantiate_from_config(
                  self.optimizer_cfg.scheduler,
                  max_decay_steps=self.trainer.max_steps,
                  lr_max=lr
              )
              scheduler = {
                  "scheduler": lr_scheduler.LambdaLR(optimizer, lr_lambda=scheduler_func.schedule),
                  "interval": "step",
                  "frequency": 1
              }
              schedulers = [scheduler]
          else:
              schedulers = []
          optimizers = [optimizer]
  
          return optimizers, schedulers
  
      @rank_zero_only
      @torch.no_grad()
      def on_train_batch_start(self, batch, batch_idx):
          # only for very first batch
          if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 \
              and batch_idx == 0 and self.ckpt_path is None:
              # set rescale weight to 1./std of encodings
              print("### USING STD-RESCALING ###")
  
              z_q = self.encode_first_stage(batch[self.first_stage_key])
              z = z_q.detach()
  
              del self.z_scale_factor
              self.register_buffer("z_scale_factor", 1. / z.flatten().std())
              print(f"setting self.z_scale_factor to {self.z_scale_factor}")
  
              print("### USING STD-RESCALING ###")
  
      def on_train_batch_end(self, *args, **kwargs):
          if self.ema_config is not None:
              self.model_ema(self.model)
  
      def on_train_epoch_start(self) -> None:
          pl.seed_everything(self.trainer.global_rank)
  
      def forward(self, batch):
          with torch.autocast(device_type="cuda", dtype=torch.bfloat16): #float32 for text
              contexts = self.cond_stage_model(image=batch.get('image'), text=batch.get('text'), mask=batch.get('mask'))
              # t5_text = contexts['t5_text']['prompt_embeds']
              # nan_count = torch.isnan(t5_text).sum()
              # if nan_count > 0:
              #     print("t5_text has %d NaN values"%(nan_count))
          with torch.autocast(device_type="cuda", dtype=torch.float16):
              with torch.no_grad():
                  latents = self.first_stage_model.encode(batch[self.first_stage_key], sample_posterior=True)
                  latents = self.z_scale_factor * latents
                  # print(latents.shape)
  
                  # check vae encode and decode is ok? answer is ok !
                  # import time
                  # from hy3dshape.pipelines import export_to_trimesh
                  # latents = 1. / self.z_scale_factor * latents
                  # latents = self.first_stage_model(latents)
                  # outputs = self.first_stage_model.latents2mesh(
                  #     latents,
                  #     bounds=1.01,
                  #     mc_level=0.0,
                  #     num_chunks=20000,
                  #     octree_resolution=256,
                  #     mc_algo='mc',
                  #     enable_pbar=True
                  # )
                  # mesh = export_to_trimesh(outputs)
                  # if isinstance(mesh, list):
                  #     for midx, m in enumerate(mesh):
                  #         m.export(f"check_{midx}_{time.time()}.glb")
                  # else:
                  #     mesh.export(f"check_{time.time()}.glb")
                  
          with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
              loss = self.transport.training_losses(self.model, latents, dict(contexts=contexts))["loss"].mean()
          return loss
  
      def training_step(self, batch, batch_idx, optimizer_idx=0):
          loss = self.forward(batch)
          split = 'train'
          loss_dict = {
              f"{split}/simple": loss.detach(),
              f"{split}/total_loss": loss.detach(),
              f"{split}/lr_abs": self.optimizers().param_groups[0]['lr'],
          }
          self.log_dict(loss_dict, prog_bar=True, logger=True, sync_dist=False, rank_zero_only=True)
  
          return loss
  
      def validation_step(self, batch, batch_idx, optimizer_idx=0):
          loss = self.forward(batch)
          split = 'val'
          loss_dict = {
              f"{split}/simple": loss.detach(),
              f"{split}/total_loss": loss.detach(),
              f"{split}/lr_abs": self.optimizers().param_groups[0]['lr'],
          }
          self.log_dict(loss_dict, prog_bar=True, logger=True, sync_dist=False, rank_zero_only=True)
  
          return loss
  
      @torch.no_grad()
      def sample(self, batch, output_type='trimesh', **kwargs):
          self.cond_stage_model.disable_drop = True
  
          generator = torch.Generator().manual_seed(0)
  
          with self.ema_scope("Sample"):
              with torch.amp.autocast(device_type='cuda'):
                  try:
                      self.pipeline.device = self.device
                      self.pipeline.dtype = self.dtype
                      print("### USING PIPELINE ###")
                      print(f'device: {self.device} dtype : {self.dtype}')
                      additional_params = {'output_type':output_type}
  
                      image = batch.get("image", None)
                      mask = batch.get('mask', None)
                      
                      # if not isinstance(image, torch.Tensor): print(image.shape)
                      # if isinstance(mask, torch.Tensor): print(mask.shape)
                      
                      outputs = self.pipeline(image=image, 
                                              mask=mask,
                                              generator=generator,
                                              **additional_params)
  
                  except Exception as e:
                      import traceback
                      traceback.print_exc()
                      print(f"Unexpected {e=}, {type(e)=}")
                      with open("error.txt", "a") as f:
                          f.write(str(e))
                          f.write(traceback.format_exc())
                          f.write("\n")
                      outputs = [None]
          self.cond_stage_model.disable_drop = False
          return [outputs]
  