content: import inspect
  import math
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  
  import comfy.utils
  
  import numpy as np
  import PIL
  import PIL.Image
  import torch
  import trimesh
  from diffusers.image_processor import PipelineImageInput
  from diffusers.pipelines.pipeline_utils import DiffusionPipeline
  from diffusers.schedulers import FlowMatchEulerDiscreteScheduler  # not sure
  from diffusers.utils import logging
  from diffusers.utils.torch_utils import randn_tensor
  from transformers import (
      BitImageProcessor,
      CLIPTokenizer,
      CLIPTextModelWithProjection,
      Dinov2Model,
  )
  from ..inference_utils import hierarchical_extract_geometry, flash_extract_geometry
  
  from ..models.autoencoders import TripoSGVAEModel
  from ..models.transformers import TripoSGDiTModel
  from .pipeline_triposg_output import TripoSGPipelineOutput
  from .pipeline_utils import TransformerDiffusionMixin
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
  def retrieve_timesteps(
      scheduler,
      num_inference_steps: Optional[int] = None,
      device: Optional[Union[str, torch.device]] = None,
      timesteps: Optional[List[int]] = None,
      sigmas: Optional[List[float]] = None,
      **kwargs,
  ):
      """
      Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
      custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
  
      Args:
          scheduler (`SchedulerMixin`):
              The scheduler to get timesteps from.
          num_inference_steps (`int`):
              The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
              must be `None`.
          device (`str` or `torch.device`, *optional*):
              The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
          timesteps (`List[int]`, *optional*):
              Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
              `num_inference_steps` and `sigmas` must be `None`.
          sigmas (`List[float]`, *optional*):
              Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
              `num_inference_steps` and `timesteps` must be `None`.
  
      Returns:
          `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
          second element is the number of inference steps.
      """
      if timesteps is not None and sigmas is not None:
          raise ValueError(
              "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
          )
      if timesteps is not None:
          accepts_timesteps = "timesteps" in set(
              inspect.signature(scheduler.set_timesteps).parameters.keys()
          )
          if not accepts_timesteps:
              raise ValueError(
                  f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                  f" timestep schedules. Please check whether you are using the correct scheduler."
              )
          scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
          timesteps = scheduler.timesteps
          num_inference_steps = len(timesteps)
      elif sigmas is not None:
          accept_sigmas = "sigmas" in set(
              inspect.signature(scheduler.set_timesteps).parameters.keys()
          )
          if not accept_sigmas:
              raise ValueError(
                  f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                  f" sigmas schedules. Please check whether you are using the correct scheduler."
              )
          scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
          timesteps = scheduler.timesteps
          num_inference_steps = len(timesteps)
      else:
          scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
          timesteps = scheduler.timesteps
      return timesteps, num_inference_steps
  
  
  class TripoSGScribblePipeline(DiffusionPipeline, TransformerDiffusionMixin):
      """
      Pipeline for (scribble and text)-to-3D generation.   
      """
  
      def __init__(
          self,
          vae: TripoSGVAEModel,
          transformer: TripoSGDiTModel,
          scheduler: FlowMatchEulerDiscreteScheduler,
          tokenizer: CLIPTokenizer,
          text_encoder: CLIPTextModelWithProjection,
          image_encoder_dinov2: Dinov2Model,
          feature_extractor_dinov2: BitImageProcessor,
      ):
          super().__init__()
  
          self.register_modules(
              vae=vae,
              transformer=transformer,
              scheduler=scheduler,
              tokenizer=tokenizer,
              text_encoder=text_encoder,
              image_encoder_dinov2=image_encoder_dinov2,
              feature_extractor_dinov2=feature_extractor_dinov2
          )
  
      @property
      def guidance_scale(self):
          return self._guidance_scale
  
      @property
      def do_classifier_free_guidance(self):
          return self._guidance_scale > 1
  
      @property
      def num_timesteps(self):
          return self._num_timesteps
  
      @property
      def attention_kwargs(self):
          return self._attention_kwargs
      
      def encode_text(self, prompt, device, num_shapes_per_prompt):
          input_ids = self.tokenizer(
              [prompt], max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
          )['input_ids'].to(device)
          text_embeds = self.text_encoder(input_ids).last_hidden_state
          text_embeds = text_embeds.repeat_interleave(num_shapes_per_prompt, dim=0)
          uncond_text_embeds = torch.zeros_like(text_embeds)
          return text_embeds, uncond_text_embeds
  
      def encode_image(self, image, device, num_shapes_per_prompt):
          dtype = next(self.image_encoder_dinov2.parameters()).dtype
  
          if not isinstance(image, torch.Tensor):
              image = self.feature_extractor_dinov2(image, return_tensors="pt").pixel_values
  
          image = image.to(device=device, dtype=dtype)
          image_embeds = self.image_encoder_dinov2(image).last_hidden_state
          image_embeds = image_embeds.repeat_interleave(num_shapes_per_prompt, dim=0)
          uncond_image_embeds = torch.zeros_like(image_embeds)
  
          return image_embeds, uncond_image_embeds
  
      def prepare_latents(
          self,
          batch_size,
          num_tokens,
          num_channels_latents,
          dtype,
          device,
          generator,
          latents: Optional[torch.Tensor] = None,
      ):
          if latents is not None:
              return latents.to(device=device, dtype=dtype)
          
          shape = (batch_size, num_tokens, num_channels_latents)
  
          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."
              )
  
          latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
  
          return latents
  
      @torch.no_grad()
      def __call__(
          self,
          image: PipelineImageInput,
          prompt: str,
          num_tokens: int = 512,
          num_inference_steps: int = 50,
          timesteps: List[int] = None,
          guidance_scale: float = 7.0,
          num_shapes_per_prompt: int = 1,
          generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
          latents: Optional[torch.FloatTensor] = None,
          attention_kwargs: Optional[Dict[str, Any]] = None,
          callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
          callback_on_step_end_tensor_inputs: List[str] = ["latents"],
          bounds: Union[Tuple[float], List[float], float] = (-1.005, -1.005, -1.005, 1.005, 1.005, 1.005),
          dense_octree_depth: int = 8, 
          hierarchical_octree_depth: int = 9,
          flash_octree_depth: int = 9,
          use_flash_decoder: bool = True,
          return_dict: bool = True,
      ):
          self._guidance_scale = guidance_scale
          self._attention_kwargs = attention_kwargs
  
          # 2. Define call parameters
          if isinstance(image, PIL.Image.Image):
              batch_size = 1
          elif isinstance(image, list):
              batch_size = len(image)
          elif isinstance(image, torch.Tensor):
              batch_size = image.shape[0]
          else:
              raise ValueError("Invalid input type for image")
  
          device = self._execution_device
  
          # 3. Encode condition
          text_embeds, negative_text_embeds = self.encode_text(
              prompt, device, num_shapes_per_prompt
          )
  
          image_embeds, negative_image_embeds = self.encode_image(
              image, device, num_shapes_per_prompt
          )        
  
          if self.do_classifier_free_guidance:
              text_embeds = torch.cat([negative_text_embeds, text_embeds], dim=0)
              image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)
  
          # 4. Prepare timesteps
          timesteps, num_inference_steps = retrieve_timesteps(
              self.scheduler, num_inference_steps, device, timesteps
          )
          num_warmup_steps = max(
              len(timesteps) - num_inference_steps * self.scheduler.order, 0
          )
          self._num_timesteps = len(timesteps)
  
          # 5. Prepare latent variables
          num_channels_latents = self.transformer.config.in_channels
          latents = self.prepare_latents(
              batch_size * num_shapes_per_prompt,
              num_tokens,
              num_channels_latents,
              image_embeds.dtype,
              device,
              generator,
              latents,
          )
  
          # 6. Denoising loop
          comfy_pbar = comfy.utils.ProgressBar(num_inference_steps)
          with self.progress_bar(total=num_inference_steps) as progress_bar:
              for i, t in enumerate(timesteps):
  
                  # 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
                  )
                  # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                  timestep = t.expand(latent_model_input.shape[0])
  
                  noise_pred = self.transformer(
                      latent_model_input,
                      timestep.to(dtype=latents.dtype),
                      encoder_hidden_states=text_embeds,
                      encoder_hidden_states_2=image_embeds,
                      attention_kwargs=attention_kwargs,
                      return_dict=False,
                  )[0]
  
                  # perform guidance
                  if self.do_classifier_free_guidance:
                      noise_pred_uncond, noise_pred_image = noise_pred.chunk(2)
                      noise_pred = noise_pred_uncond + self.guidance_scale * (
                          noise_pred_image - noise_pred_uncond
                      )
  
                  # compute the previous noisy sample x_t -> x_t-1
                  latents_dtype = latents.dtype
                  latents = self.scheduler.step(
                      noise_pred, t, latents, 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)
  
                  # 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()
                      
                  comfy_pbar.update_absolute(i + 1)
  
          # 7. decoder mesh
          if not use_flash_decoder:
              geometric_func = lambda x: self.vae.decode(latents, sampled_points=x).sample
              output = hierarchical_extract_geometry(
                  geometric_func,
                  device,
                  bounds=bounds,
                  dense_octree_depth=dense_octree_depth,
                  hierarchical_octree_depth=hierarchical_octree_depth,
              )
          else:
              self.vae.set_flash_decoder()
              output = flash_extract_geometry(
                  latents,
                  self.vae,
                  bounds=bounds,
                  octree_depth=flash_octree_depth,
              )
          meshes = [trimesh.Trimesh(mesh_v_f[0].astype(np.float32), mesh_v_f[1]) for mesh_v_f in output]
          
          # Offload all models
          self.maybe_free_model_hooks()
  
          if not return_dict:
              return (output, meshes)
  
          return TripoSGPipelineOutput(samples=output, meshes=meshes)
  