SD 2.1. Mari

SD2JointModel [nn.Module]

config:
  - prediction_type <- prediction_type = "v_prediction"
  - use_ema <- use_ema = True
  - freeze_encoder <- freeze_encoder = False
  - freeze_cross_attention <- freeze_cross_attention = False
  - use_deterministic_vae <- use_deterministic_vae = False
  - lambda_dehaze <- lambda_dehaze = 1.0
  - lambda_depth <- lambda_depth = 1.0
  - diffusion_prediction_loss <- diffusion_prediction_loss = None
  ... +113 more

submodules:
  - vae (AutoencoderKL.from_pretrained)
  - noise_scheduler (DDPMScheduler.from_pretrained)
  - inference_scheduler (DDIMScheduler.from_pretrained)
  - conv_in (unet.conv_in)
  - time_proj (unet.time_proj)
  - time_embedding (unet.time_embedding)
  - down_blocks (unet.down_blocks)
  - mid_block (unet.mid_block)
  - dehaze_up_blocks (unet.up_blocks)
  - dehaze_conv_norm_out (unet.conv_norm_out)
  ... +18 more

methods:
  - forward(hazy_images, clear_images, depth_images, valid_mask, timesteps, visibility_loss_scale, return_camera, return_geometry, external_intrinsics, external_K, camera_valid, rays_homogeneous, rays_unit, ray_valid_mask, sky_mask, sky_sdf, sky_label_weight, training_epoch) -> _sample_training_timesteps, forward_joint, forward_dehaze
  - forward_dehaze(noisy_clear, hazy_latents, timesteps, encoder_hidden_states, vis_ctx, feature_ctx) -> _time_embed, _encoder_forward, _collect_depth_guidance
  - forward_depth(noisy_depth, hazy_latents, timesteps, encoder_hidden_states, vis_ctx, feature_ctx) -> _time_embed, _encoder_forward, _collect_dehaze_guidance
  - forward_joint(noisy_rgb, noisy_depth, hazy_latents, timesteps, encoder_hidden_states, vis_ctx, noisy_segmentation, _task_order, feature_ctx) -> _time_embed, _encoder_forward, _joint_decoder_blocks_forward
  - _decoder_blocks_forward(sample, down_block_res_samples, emb, encoder_hidden_states, up_blocks, up_adapters, vis_ctx, fusion_vis_ctx, collect_indices, stop_after_block, fusion_features, fusion_adapters, skip_adapters)
  - _decoder_forward(sample, down_block_res_samples, emb, encoder_hidden_states, up_blocks, conv_norm_out, conv_out, up_adapters, out_adapter, vis_ctx, fusion_vis_ctx, fusion_features, fusion_adapters, skip_adapters) -> _decoder_blocks_forward
  - _encoder_forward(sample, emb, encoder_hidden_states, feature_ctx)
  - _joint_decoder_blocks_forward(sample, skips, emb, encoder_hidden_states, dehaze_vis_ctx, depth_vis_ctx, fusion_vis_ctx, _task_order)
  - _collect_dehaze_guidance(sample, down_block_res_samples, emb, encoder_hidden_states, vis_ctx) -> _decoder_blocks_forward
  - _collect_depth_guidance(sample, down_block_res_samples, emb, encoder_hidden_states, vis_ctx) -> _decoder_blocks_forward
  ... +3 more