Copy disabled (too large)
Download .txt
Showing preview only (10,421K chars total). Download the full file to get everything.
Repository: kxhit/EscherNet
Branch: main
Commit: 10b650492ba9
Files: 768
Total size: 126.9 MB
Directory structure:
gitextract_grlqwy99/
├── .gitignore
├── 3drecon/
│ ├── configs/
│ │ └── neus_36.yaml
│ ├── ours_GSO_T1/
│ │ └── NeuS/
│ │ ├── grandmother/
│ │ │ ├── mesh.ply
│ │ │ └── tensorboard_logs/
│ │ │ └── version_0/
│ │ │ ├── events.out.tfevents.1706660493.xin-PC.934205.0
│ │ │ └── hparams.yaml
│ │ ├── lion/
│ │ │ ├── mesh.ply
│ │ │ └── tensorboard_logs/
│ │ │ └── version_0/
│ │ │ ├── events.out.tfevents.1706659999.xin-PC.931603.0
│ │ │ └── hparams.yaml
│ │ └── train/
│ │ └── tensorboard_logs/
│ │ └── version_0/
│ │ ├── events.out.tfevents.1706660976.xin-PC.936418.0
│ │ └── hparams.yaml
│ ├── raymarching/
│ │ ├── __init__.py
│ │ ├── backend.py
│ │ ├── raymarching.py
│ │ ├── setup.py
│ │ └── src/
│ │ ├── bindings.cpp
│ │ ├── raymarching.cu
│ │ └── raymarching.h
│ ├── renderer/
│ │ ├── agg_net.py
│ │ ├── cost_reg_net.py
│ │ ├── dummy_dataset.py
│ │ ├── feature_net.py
│ │ ├── neus_networks.py
│ │ ├── ngp_renderer.py
│ │ └── renderer.py
│ ├── run_NeuS.py
│ ├── train_renderer.py
│ └── util.py
├── 4DoF/
│ ├── CN_encoder.py
│ ├── dataset.py
│ ├── diffusers/
│ │ ├── __init__.py
│ │ ├── commands/
│ │ │ ├── __init__.py
│ │ │ ├── diffusers_cli.py
│ │ │ └── env.py
│ │ ├── configuration_utils.py
│ │ ├── dependency_versions_check.py
│ │ ├── dependency_versions_table.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ └── rl/
│ │ │ ├── __init__.py
│ │ │ └── value_guided_sampling.py
│ │ ├── image_processor.py
│ │ ├── loaders.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ ├── activations.py
│ │ │ ├── attention.py
│ │ │ ├── attention_flax.py
│ │ │ ├── attention_processor.py
│ │ │ ├── autoencoder_kl.py
│ │ │ ├── controlnet.py
│ │ │ ├── controlnet_flax.py
│ │ │ ├── cross_attention.py
│ │ │ ├── dual_transformer_2d.py
│ │ │ ├── embeddings.py
│ │ │ ├── embeddings_flax.py
│ │ │ ├── modeling_flax_pytorch_utils.py
│ │ │ ├── modeling_flax_utils.py
│ │ │ ├── modeling_pytorch_flax_utils.py
│ │ │ ├── modeling_utils.py
│ │ │ ├── prior_transformer.py
│ │ │ ├── resnet.py
│ │ │ ├── resnet_flax.py
│ │ │ ├── t5_film_transformer.py
│ │ │ ├── transformer_2d.py
│ │ │ ├── transformer_temporal.py
│ │ │ ├── unet_1d.py
│ │ │ ├── unet_1d_blocks.py
│ │ │ ├── unet_2d.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ ├── unet_2d_blocks_flax.py
│ │ │ ├── unet_2d_condition.py
│ │ │ ├── unet_2d_condition_flax.py
│ │ │ ├── unet_3d_blocks.py
│ │ │ ├── unet_3d_condition.py
│ │ │ ├── vae.py
│ │ │ ├── vae_flax.py
│ │ │ └── vq_model.py
│ │ ├── optimization.py
│ │ ├── pipeline_utils.py
│ │ ├── pipelines/
│ │ │ ├── __init__.py
│ │ │ ├── alt_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_roberta_series.py
│ │ │ │ ├── pipeline_alt_diffusion.py
│ │ │ │ └── pipeline_alt_diffusion_img2img.py
│ │ │ ├── audio_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mel.py
│ │ │ │ └── pipeline_audio_diffusion.py
│ │ │ ├── audioldm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_audioldm.py
│ │ │ ├── consistency_models/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_consistency_models.py
│ │ │ ├── controlnet/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── multicontrolnet.py
│ │ │ │ ├── pipeline_controlnet.py
│ │ │ │ ├── pipeline_controlnet_img2img.py
│ │ │ │ ├── pipeline_controlnet_inpaint.py
│ │ │ │ └── pipeline_flax_controlnet.py
│ │ │ ├── dance_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dance_diffusion.py
│ │ │ ├── ddim/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddim.py
│ │ │ ├── ddpm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddpm.py
│ │ │ ├── deepfloyd_if/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_if.py
│ │ │ │ ├── pipeline_if_img2img.py
│ │ │ │ ├── pipeline_if_img2img_superresolution.py
│ │ │ │ ├── pipeline_if_inpainting.py
│ │ │ │ ├── pipeline_if_inpainting_superresolution.py
│ │ │ │ ├── pipeline_if_superresolution.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── timesteps.py
│ │ │ │ └── watermark.py
│ │ │ ├── dit/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dit.py
│ │ │ ├── kandinsky/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky.py
│ │ │ │ ├── pipeline_kandinsky_img2img.py
│ │ │ │ ├── pipeline_kandinsky_inpaint.py
│ │ │ │ ├── pipeline_kandinsky_prior.py
│ │ │ │ └── text_encoder.py
│ │ │ ├── kandinsky2_2/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky2_2.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_inpainting.py
│ │ │ │ ├── pipeline_kandinsky2_2_prior.py
│ │ │ │ └── pipeline_kandinsky2_2_prior_emb2emb.py
│ │ │ ├── latent_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_latent_diffusion.py
│ │ │ │ └── pipeline_latent_diffusion_superresolution.py
│ │ │ ├── latent_diffusion_uncond/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_latent_diffusion_uncond.py
│ │ │ ├── onnx_utils.py
│ │ │ ├── paint_by_example/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── pipeline_paint_by_example.py
│ │ │ ├── pipeline_flax_utils.py
│ │ │ ├── pipeline_utils.py
│ │ │ ├── pndm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_pndm.py
│ │ │ ├── repaint/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_repaint.py
│ │ │ ├── score_sde_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_score_sde_ve.py
│ │ │ ├── semantic_stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_semantic_stable_diffusion.py
│ │ │ ├── shap_e/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── camera.py
│ │ │ │ ├── pipeline_shap_e.py
│ │ │ │ ├── pipeline_shap_e_img2img.py
│ │ │ │ └── renderer.py
│ │ │ ├── spectrogram_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── continous_encoder.py
│ │ │ │ ├── midi_utils.py
│ │ │ │ ├── notes_encoder.py
│ │ │ │ └── pipeline_spectrogram_diffusion.py
│ │ │ ├── stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── convert_from_ckpt.py
│ │ │ │ ├── pipeline_cycle_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_attend_and_excite.py
│ │ │ │ ├── pipeline_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_stable_diffusion_depth2img.py
│ │ │ │ ├── pipeline_stable_diffusion_diffedit.py
│ │ │ │ ├── pipeline_stable_diffusion_image_variation.py
│ │ │ │ ├── pipeline_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_stable_diffusion_instruct_pix2pix.py
│ │ │ │ ├── pipeline_stable_diffusion_k_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_latent_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion_ldm3d.py
│ │ │ │ ├── pipeline_stable_diffusion_model_editing.py
│ │ │ │ ├── pipeline_stable_diffusion_panorama.py
│ │ │ │ ├── pipeline_stable_diffusion_paradigms.py
│ │ │ │ ├── pipeline_stable_diffusion_pix2pix_zero.py
│ │ │ │ ├── pipeline_stable_diffusion_sag.py
│ │ │ │ ├── pipeline_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_unclip.py
│ │ │ │ ├── pipeline_stable_unclip_img2img.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── safety_checker_flax.py
│ │ │ │ └── stable_unclip_image_normalizer.py
│ │ │ ├── stable_diffusion_safe/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_safe.py
│ │ │ │ └── safety_checker.py
│ │ │ ├── stable_diffusion_xl/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_xl.py
│ │ │ │ ├── pipeline_stable_diffusion_xl_img2img.py
│ │ │ │ └── watermark.py
│ │ │ ├── stochastic_karras_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_stochastic_karras_ve.py
│ │ │ ├── text_to_video_synthesis/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_text_to_video_synth.py
│ │ │ │ ├── pipeline_text_to_video_synth_img2img.py
│ │ │ │ └── pipeline_text_to_video_zero.py
│ │ │ ├── unclip/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_unclip.py
│ │ │ │ ├── pipeline_unclip_image_variation.py
│ │ │ │ └── text_proj.py
│ │ │ ├── unidiffuser/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_decoder.py
│ │ │ │ ├── modeling_uvit.py
│ │ │ │ └── pipeline_unidiffuser.py
│ │ │ ├── versatile_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_unet.py
│ │ │ │ ├── pipeline_versatile_diffusion.py
│ │ │ │ ├── pipeline_versatile_diffusion_dual_guided.py
│ │ │ │ ├── pipeline_versatile_diffusion_image_variation.py
│ │ │ │ └── pipeline_versatile_diffusion_text_to_image.py
│ │ │ └── vq_diffusion/
│ │ │ ├── __init__.py
│ │ │ └── pipeline_vq_diffusion.py
│ │ ├── schedulers/
│ │ │ ├── __init__.py
│ │ │ ├── scheduling_consistency_models.py
│ │ │ ├── scheduling_ddim.py
│ │ │ ├── scheduling_ddim_flax.py
│ │ │ ├── scheduling_ddim_inverse.py
│ │ │ ├── scheduling_ddim_parallel.py
│ │ │ ├── scheduling_ddpm.py
│ │ │ ├── scheduling_ddpm_flax.py
│ │ │ ├── scheduling_ddpm_parallel.py
│ │ │ ├── scheduling_deis_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep_flax.py
│ │ │ ├── scheduling_dpmsolver_multistep_inverse.py
│ │ │ ├── scheduling_dpmsolver_sde.py
│ │ │ ├── scheduling_dpmsolver_singlestep.py
│ │ │ ├── scheduling_euler_ancestral_discrete.py
│ │ │ ├── scheduling_euler_discrete.py
│ │ │ ├── scheduling_heun_discrete.py
│ │ │ ├── scheduling_ipndm.py
│ │ │ ├── scheduling_k_dpm_2_ancestral_discrete.py
│ │ │ ├── scheduling_k_dpm_2_discrete.py
│ │ │ ├── scheduling_karras_ve.py
│ │ │ ├── scheduling_karras_ve_flax.py
│ │ │ ├── scheduling_lms_discrete.py
│ │ │ ├── scheduling_lms_discrete_flax.py
│ │ │ ├── scheduling_pndm.py
│ │ │ ├── scheduling_pndm_flax.py
│ │ │ ├── scheduling_repaint.py
│ │ │ ├── scheduling_sde_ve.py
│ │ │ ├── scheduling_sde_ve_flax.py
│ │ │ ├── scheduling_sde_vp.py
│ │ │ ├── scheduling_unclip.py
│ │ │ ├── scheduling_unipc_multistep.py
│ │ │ ├── scheduling_utils.py
│ │ │ ├── scheduling_utils_flax.py
│ │ │ └── scheduling_vq_diffusion.py
│ │ ├── training_utils.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── accelerate_utils.py
│ │ ├── constants.py
│ │ ├── deprecation_utils.py
│ │ ├── doc_utils.py
│ │ ├── dummy_flax_and_transformers_objects.py
│ │ ├── dummy_flax_objects.py
│ │ ├── dummy_note_seq_objects.py
│ │ ├── dummy_onnx_objects.py
│ │ ├── dummy_pt_objects.py
│ │ ├── dummy_torch_and_librosa_objects.py
│ │ ├── dummy_torch_and_scipy_objects.py
│ │ ├── dummy_torch_and_torchsde_objects.py
│ │ ├── dummy_torch_and_transformers_and_invisible_watermark_objects.py
│ │ ├── dummy_torch_and_transformers_and_k_diffusion_objects.py
│ │ ├── dummy_torch_and_transformers_and_onnx_objects.py
│ │ ├── dummy_torch_and_transformers_objects.py
│ │ ├── dummy_transformers_and_torch_and_note_seq_objects.py
│ │ ├── dynamic_modules_utils.py
│ │ ├── hub_utils.py
│ │ ├── import_utils.py
│ │ ├── logging.py
│ │ ├── model_card_template.md
│ │ ├── outputs.py
│ │ ├── pil_utils.py
│ │ ├── testing_utils.py
│ │ └── torch_utils.py
│ ├── pipeline_zero1to3.py
│ ├── train_eschernet.py
│ └── unet_2d_condition.py
├── 6DoF/
│ ├── CN_encoder.py
│ ├── dataset.py
│ ├── diffusers/
│ │ ├── __init__.py
│ │ ├── commands/
│ │ │ ├── __init__.py
│ │ │ ├── diffusers_cli.py
│ │ │ └── env.py
│ │ ├── configuration_utils.py
│ │ ├── dependency_versions_check.py
│ │ ├── dependency_versions_table.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ └── rl/
│ │ │ ├── __init__.py
│ │ │ └── value_guided_sampling.py
│ │ ├── image_processor.py
│ │ ├── loaders.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ ├── activations.py
│ │ │ ├── attention.py
│ │ │ ├── attention_flax.py
│ │ │ ├── attention_processor.py
│ │ │ ├── autoencoder_kl.py
│ │ │ ├── controlnet.py
│ │ │ ├── controlnet_flax.py
│ │ │ ├── cross_attention.py
│ │ │ ├── dual_transformer_2d.py
│ │ │ ├── embeddings.py
│ │ │ ├── embeddings_flax.py
│ │ │ ├── modeling_flax_pytorch_utils.py
│ │ │ ├── modeling_flax_utils.py
│ │ │ ├── modeling_pytorch_flax_utils.py
│ │ │ ├── modeling_utils.py
│ │ │ ├── prior_transformer.py
│ │ │ ├── resnet.py
│ │ │ ├── resnet_flax.py
│ │ │ ├── t5_film_transformer.py
│ │ │ ├── transformer_2d.py
│ │ │ ├── transformer_temporal.py
│ │ │ ├── unet_1d.py
│ │ │ ├── unet_1d_blocks.py
│ │ │ ├── unet_2d.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ ├── unet_2d_blocks_flax.py
│ │ │ ├── unet_2d_condition.py
│ │ │ ├── unet_2d_condition_flax.py
│ │ │ ├── unet_3d_blocks.py
│ │ │ ├── unet_3d_condition.py
│ │ │ ├── vae.py
│ │ │ ├── vae_flax.py
│ │ │ └── vq_model.py
│ │ ├── optimization.py
│ │ ├── pipeline_utils.py
│ │ ├── pipelines/
│ │ │ ├── __init__.py
│ │ │ ├── alt_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_roberta_series.py
│ │ │ │ ├── pipeline_alt_diffusion.py
│ │ │ │ └── pipeline_alt_diffusion_img2img.py
│ │ │ ├── audio_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mel.py
│ │ │ │ └── pipeline_audio_diffusion.py
│ │ │ ├── audioldm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_audioldm.py
│ │ │ ├── consistency_models/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_consistency_models.py
│ │ │ ├── controlnet/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── multicontrolnet.py
│ │ │ │ ├── pipeline_controlnet.py
│ │ │ │ ├── pipeline_controlnet_img2img.py
│ │ │ │ ├── pipeline_controlnet_inpaint.py
│ │ │ │ └── pipeline_flax_controlnet.py
│ │ │ ├── dance_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dance_diffusion.py
│ │ │ ├── ddim/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddim.py
│ │ │ ├── ddpm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddpm.py
│ │ │ ├── deepfloyd_if/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_if.py
│ │ │ │ ├── pipeline_if_img2img.py
│ │ │ │ ├── pipeline_if_img2img_superresolution.py
│ │ │ │ ├── pipeline_if_inpainting.py
│ │ │ │ ├── pipeline_if_inpainting_superresolution.py
│ │ │ │ ├── pipeline_if_superresolution.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── timesteps.py
│ │ │ │ └── watermark.py
│ │ │ ├── dit/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dit.py
│ │ │ ├── kandinsky/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky.py
│ │ │ │ ├── pipeline_kandinsky_img2img.py
│ │ │ │ ├── pipeline_kandinsky_inpaint.py
│ │ │ │ ├── pipeline_kandinsky_prior.py
│ │ │ │ └── text_encoder.py
│ │ │ ├── kandinsky2_2/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky2_2.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_inpainting.py
│ │ │ │ ├── pipeline_kandinsky2_2_prior.py
│ │ │ │ └── pipeline_kandinsky2_2_prior_emb2emb.py
│ │ │ ├── latent_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_latent_diffusion.py
│ │ │ │ └── pipeline_latent_diffusion_superresolution.py
│ │ │ ├── latent_diffusion_uncond/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_latent_diffusion_uncond.py
│ │ │ ├── onnx_utils.py
│ │ │ ├── paint_by_example/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── pipeline_paint_by_example.py
│ │ │ ├── pipeline_flax_utils.py
│ │ │ ├── pipeline_utils.py
│ │ │ ├── pndm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_pndm.py
│ │ │ ├── repaint/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_repaint.py
│ │ │ ├── score_sde_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_score_sde_ve.py
│ │ │ ├── semantic_stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_semantic_stable_diffusion.py
│ │ │ ├── shap_e/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── camera.py
│ │ │ │ ├── pipeline_shap_e.py
│ │ │ │ ├── pipeline_shap_e_img2img.py
│ │ │ │ └── renderer.py
│ │ │ ├── spectrogram_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── continous_encoder.py
│ │ │ │ ├── midi_utils.py
│ │ │ │ ├── notes_encoder.py
│ │ │ │ └── pipeline_spectrogram_diffusion.py
│ │ │ ├── stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── convert_from_ckpt.py
│ │ │ │ ├── pipeline_cycle_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_attend_and_excite.py
│ │ │ │ ├── pipeline_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_stable_diffusion_depth2img.py
│ │ │ │ ├── pipeline_stable_diffusion_diffedit.py
│ │ │ │ ├── pipeline_stable_diffusion_image_variation.py
│ │ │ │ ├── pipeline_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_stable_diffusion_instruct_pix2pix.py
│ │ │ │ ├── pipeline_stable_diffusion_k_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_latent_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion_ldm3d.py
│ │ │ │ ├── pipeline_stable_diffusion_model_editing.py
│ │ │ │ ├── pipeline_stable_diffusion_panorama.py
│ │ │ │ ├── pipeline_stable_diffusion_paradigms.py
│ │ │ │ ├── pipeline_stable_diffusion_pix2pix_zero.py
│ │ │ │ ├── pipeline_stable_diffusion_sag.py
│ │ │ │ ├── pipeline_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_unclip.py
│ │ │ │ ├── pipeline_stable_unclip_img2img.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── safety_checker_flax.py
│ │ │ │ └── stable_unclip_image_normalizer.py
│ │ │ ├── stable_diffusion_safe/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_safe.py
│ │ │ │ └── safety_checker.py
│ │ │ ├── stable_diffusion_xl/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_xl.py
│ │ │ │ ├── pipeline_stable_diffusion_xl_img2img.py
│ │ │ │ └── watermark.py
│ │ │ ├── stochastic_karras_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_stochastic_karras_ve.py
│ │ │ ├── text_to_video_synthesis/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_text_to_video_synth.py
│ │ │ │ ├── pipeline_text_to_video_synth_img2img.py
│ │ │ │ └── pipeline_text_to_video_zero.py
│ │ │ ├── unclip/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_unclip.py
│ │ │ │ ├── pipeline_unclip_image_variation.py
│ │ │ │ └── text_proj.py
│ │ │ ├── unidiffuser/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_decoder.py
│ │ │ │ ├── modeling_uvit.py
│ │ │ │ └── pipeline_unidiffuser.py
│ │ │ ├── versatile_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_unet.py
│ │ │ │ ├── pipeline_versatile_diffusion.py
│ │ │ │ ├── pipeline_versatile_diffusion_dual_guided.py
│ │ │ │ ├── pipeline_versatile_diffusion_image_variation.py
│ │ │ │ └── pipeline_versatile_diffusion_text_to_image.py
│ │ │ └── vq_diffusion/
│ │ │ ├── __init__.py
│ │ │ └── pipeline_vq_diffusion.py
│ │ ├── schedulers/
│ │ │ ├── __init__.py
│ │ │ ├── scheduling_consistency_models.py
│ │ │ ├── scheduling_ddim.py
│ │ │ ├── scheduling_ddim_flax.py
│ │ │ ├── scheduling_ddim_inverse.py
│ │ │ ├── scheduling_ddim_parallel.py
│ │ │ ├── scheduling_ddpm.py
│ │ │ ├── scheduling_ddpm_flax.py
│ │ │ ├── scheduling_ddpm_parallel.py
│ │ │ ├── scheduling_deis_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep_flax.py
│ │ │ ├── scheduling_dpmsolver_multistep_inverse.py
│ │ │ ├── scheduling_dpmsolver_sde.py
│ │ │ ├── scheduling_dpmsolver_singlestep.py
│ │ │ ├── scheduling_euler_ancestral_discrete.py
│ │ │ ├── scheduling_euler_discrete.py
│ │ │ ├── scheduling_heun_discrete.py
│ │ │ ├── scheduling_ipndm.py
│ │ │ ├── scheduling_k_dpm_2_ancestral_discrete.py
│ │ │ ├── scheduling_k_dpm_2_discrete.py
│ │ │ ├── scheduling_karras_ve.py
│ │ │ ├── scheduling_karras_ve_flax.py
│ │ │ ├── scheduling_lms_discrete.py
│ │ │ ├── scheduling_lms_discrete_flax.py
│ │ │ ├── scheduling_pndm.py
│ │ │ ├── scheduling_pndm_flax.py
│ │ │ ├── scheduling_repaint.py
│ │ │ ├── scheduling_sde_ve.py
│ │ │ ├── scheduling_sde_ve_flax.py
│ │ │ ├── scheduling_sde_vp.py
│ │ │ ├── scheduling_unclip.py
│ │ │ ├── scheduling_unipc_multistep.py
│ │ │ ├── scheduling_utils.py
│ │ │ ├── scheduling_utils_flax.py
│ │ │ └── scheduling_vq_diffusion.py
│ │ ├── training_utils.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── accelerate_utils.py
│ │ ├── constants.py
│ │ ├── deprecation_utils.py
│ │ ├── doc_utils.py
│ │ ├── dummy_flax_and_transformers_objects.py
│ │ ├── dummy_flax_objects.py
│ │ ├── dummy_note_seq_objects.py
│ │ ├── dummy_onnx_objects.py
│ │ ├── dummy_pt_objects.py
│ │ ├── dummy_torch_and_librosa_objects.py
│ │ ├── dummy_torch_and_scipy_objects.py
│ │ ├── dummy_torch_and_torchsde_objects.py
│ │ ├── dummy_torch_and_transformers_and_invisible_watermark_objects.py
│ │ ├── dummy_torch_and_transformers_and_k_diffusion_objects.py
│ │ ├── dummy_torch_and_transformers_and_onnx_objects.py
│ │ ├── dummy_torch_and_transformers_objects.py
│ │ ├── dummy_transformers_and_torch_and_note_seq_objects.py
│ │ ├── dynamic_modules_utils.py
│ │ ├── hub_utils.py
│ │ ├── import_utils.py
│ │ ├── logging.py
│ │ ├── model_card_template.md
│ │ ├── outputs.py
│ │ ├── pil_utils.py
│ │ ├── testing_utils.py
│ │ └── torch_utils.py
│ ├── instructions.md
│ ├── pipeline_zero1to3.py
│ ├── train_eschernet.py
│ └── unet_2d_condition.py
├── CaPE.py
├── LICENSE
├── README.md
├── demo/
│ └── GSO30/
│ └── elephant/
│ └── render_mvs_25/
│ └── model/
│ ├── 000.npy
│ ├── 001.npy
│ ├── 002.npy
│ ├── 003.npy
│ ├── 004.npy
│ ├── 005.npy
│ ├── 006.npy
│ ├── 007.npy
│ ├── 008.npy
│ ├── 009.npy
│ ├── 010.npy
│ ├── 011.npy
│ ├── 012.npy
│ ├── 013.npy
│ ├── 014.npy
│ ├── 015.npy
│ ├── 016.npy
│ ├── 017.npy
│ ├── 018.npy
│ ├── 019.npy
│ ├── 020.npy
│ ├── 021.npy
│ ├── 022.npy
│ ├── 023.npy
│ └── 024.npy
├── environment.yml
├── eval_eschernet.py
├── eval_eschernet.sh
├── metrics/
│ ├── NeRF_idx/
│ │ ├── chair/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── drums/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── ficus/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── hotdog/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── lego/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── materials/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── mic/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ └── ship/
│ │ ├── test_M20.npy
│ │ ├── train_N100M20_random.npy
│ │ ├── train_N10M20_random.npy
│ │ ├── train_N1M20_random.npy
│ │ ├── train_N20M20_random.npy
│ │ ├── train_N2M20_random.npy
│ │ ├── train_N30M20_random.npy
│ │ ├── train_N3M20_random.npy
│ │ ├── train_N50M20_random.npy
│ │ ├── train_N5M20_random.npy
│ │ ├── train_fov0M20_fov.npy
│ │ ├── train_fov135M20_fov.npy
│ │ ├── train_fov15M20_fov.npy
│ │ ├── train_fov180M20_fov.npy
│ │ ├── train_fov225M20_fov.npy
│ │ ├── train_fov270M20_fov.npy
│ │ ├── train_fov30M20_fov.npy
│ │ ├── train_fov360M20_fov.npy
│ │ ├── train_fov50M20_fov.npy
│ │ ├── train_fov5M20_fov.npy
│ │ └── train_fov90M20_fov.npy
│ ├── RenderOption_phong.json
│ ├── RenderOption_rgb.json
│ ├── ScreenCamera_0.json
│ ├── ScreenCamera_1.json
│ ├── ScreenCamera_2.json
│ ├── ScreenCamera_3.json
│ ├── ScreenCamera_4.json
│ ├── ScreenCamera_5.json
│ ├── eval_2D_NVS.py
│ ├── eval_3D_GSO.py
│ └── metrics.py
└── scripts/
├── README.md
├── all_invalid.npy
├── blender_script_mvs.py
├── empty_ids.npy
├── invalid_ids.npy
├── obj_ids.npy
├── objaverse_filter.py
└── render_all_mvs.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
.idea/
training/
lightning_logs/
image_log/
*.pth
*.pt
*.ckpt
*.safetensors
gradio_pose2image_private.py
gradio_canny2image_private.py
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# log files
log_bigboy
logs*
wandb
# pretrained models
pretrained
================================================
FILE: 3drecon/configs/neus_36.yaml
================================================
model:
base_lr: 5.0e-4
target: renderer.renderer.RendererTrainer
params:
total_steps: 2000
warm_up_steps: 100
train_batch_num: 2560
train_batch_fg_num: 512
test_batch_num: 4096
use_mask: true
lambda_rgb_loss: 0.5
lambda_mask_loss: 1.0
lambda_eikonal_loss: 0.1
use_warm_up: true
data:
target: renderer.dummy_dataset.DummyDataset
params: {}
callbacks:
save_interval: 500
trainer:
val_check_interval: 500
max_steps: 2000
================================================
FILE: 3drecon/ours_GSO_T1/NeuS/grandmother/mesh.ply
================================================
[File too large to display: 10.4 MB]
================================================
FILE: 3drecon/ours_GSO_T1/NeuS/grandmother/tensorboard_logs/version_0/hparams.yaml
================================================
{}
================================================
FILE: 3drecon/ours_GSO_T1/NeuS/lion/mesh.ply
================================================
[File too large to display: 10.3 MB]
================================================
FILE: 3drecon/ours_GSO_T1/NeuS/lion/tensorboard_logs/version_0/hparams.yaml
================================================
{}
================================================
FILE: 3drecon/ours_GSO_T1/NeuS/train/tensorboard_logs/version_0/hparams.yaml
================================================
{}
================================================
FILE: 3drecon/raymarching/__init__.py
================================================
from .raymarching import *
================================================
FILE: 3drecon/raymarching/backend.py
================================================
import os
from torch.utils.cpp_extension import load
_src_path = os.path.dirname(os.path.abspath(__file__))
nvcc_flags = [
'-O3', '-std=c++14',
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__',
]
if os.name == "posix":
c_flags = ['-O3', '-std=c++14']
elif os.name == "nt":
c_flags = ['/O2', '/std:c++17']
# find cl.exe
def find_cl_path():
import glob
for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
paths = sorted(glob.glob(r"C:\\Program Files (x86)\\Microsoft Visual Studio\\*\\%s\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64" % edition), reverse=True)
if paths:
return paths[0]
# If cl.exe is not on path, try to find it.
if os.system("where cl.exe >nul 2>nul") != 0:
cl_path = find_cl_path()
if cl_path is None:
raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
os.environ["PATH"] += ";" + cl_path
_backend = load(name='_raymarching',
extra_cflags=c_flags,
extra_cuda_cflags=nvcc_flags,
sources=[os.path.join(_src_path, 'src', f) for f in [
'raymarching.cu',
'bindings.cpp',
]],
)
__all__ = ['_backend']
================================================
FILE: 3drecon/raymarching/raymarching.py
================================================
import numpy as np
import time
import torch
import torch.nn as nn
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd
try:
import _raymarching as _backend
except ImportError:
from .backend import _backend
# ----------------------------------------
# utils
# ----------------------------------------
class _near_far_from_aabb(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, rays_o, rays_d, aabb, min_near=0.2):
''' near_far_from_aabb, CUDA implementation
Calculate rays' intersection time (near and far) with aabb
Args:
rays_o: float, [N, 3]
rays_d: float, [N, 3]
aabb: float, [6], (xmin, ymin, zmin, xmax, ymax, zmax)
min_near: float, scalar
Returns:
nears: float, [N]
fars: float, [N]
'''
if not rays_o.is_cuda: rays_o = rays_o.cuda()
if not rays_d.is_cuda: rays_d = rays_d.cuda()
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
N = rays_o.shape[0] # num rays
nears = torch.empty(N, dtype=rays_o.dtype, device=rays_o.device)
fars = torch.empty(N, dtype=rays_o.dtype, device=rays_o.device)
_backend.near_far_from_aabb(rays_o, rays_d, aabb, N, min_near, nears, fars)
return nears, fars
near_far_from_aabb = _near_far_from_aabb.apply
class _sph_from_ray(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, rays_o, rays_d, radius):
''' sph_from_ray, CUDA implementation
get spherical coordinate on the background sphere from rays.
Assume rays_o are inside the Sphere(radius).
Args:
rays_o: [N, 3]
rays_d: [N, 3]
radius: scalar, float
Return:
coords: [N, 2], in [-1, 1], theta and phi on a sphere. (further-surface)
'''
if not rays_o.is_cuda: rays_o = rays_o.cuda()
if not rays_d.is_cuda: rays_d = rays_d.cuda()
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
N = rays_o.shape[0] # num rays
coords = torch.empty(N, 2, dtype=rays_o.dtype, device=rays_o.device)
_backend.sph_from_ray(rays_o, rays_d, radius, N, coords)
return coords
sph_from_ray = _sph_from_ray.apply
class _morton3D(Function):
@staticmethod
def forward(ctx, coords):
''' morton3D, CUDA implementation
Args:
coords: [N, 3], int32, in [0, 128) (for some reason there is no uint32 tensor in torch...)
TODO: check if the coord range is valid! (current 128 is safe)
Returns:
indices: [N], int32, in [0, 128^3)
'''
if not coords.is_cuda: coords = coords.cuda()
N = coords.shape[0]
indices = torch.empty(N, dtype=torch.int32, device=coords.device)
_backend.morton3D(coords.int(), N, indices)
return indices
morton3D = _morton3D.apply
class _morton3D_invert(Function):
@staticmethod
def forward(ctx, indices):
''' morton3D_invert, CUDA implementation
Args:
indices: [N], int32, in [0, 128^3)
Returns:
coords: [N, 3], int32, in [0, 128)
'''
if not indices.is_cuda: indices = indices.cuda()
N = indices.shape[0]
coords = torch.empty(N, 3, dtype=torch.int32, device=indices.device)
_backend.morton3D_invert(indices.int(), N, coords)
return coords
morton3D_invert = _morton3D_invert.apply
class _packbits(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, grid, thresh, bitfield=None):
''' packbits, CUDA implementation
Pack up the density grid into a bit field to accelerate ray marching.
Args:
grid: float, [C, H * H * H], assume H % 2 == 0
thresh: float, threshold
Returns:
bitfield: uint8, [C, H * H * H / 8]
'''
if not grid.is_cuda: grid = grid.cuda()
grid = grid.contiguous()
C = grid.shape[0]
H3 = grid.shape[1]
N = C * H3 // 8
if bitfield is None:
bitfield = torch.empty(N, dtype=torch.uint8, device=grid.device)
_backend.packbits(grid, N, thresh, bitfield)
return bitfield
packbits = _packbits.apply
# ----------------------------------------
# train functions
# ----------------------------------------
class _march_rays_train(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, rays_o, rays_d, bound, density_bitfield, C, H, nears, fars, step_counter=None, mean_count=-1, perturb=False, align=-1, force_all_rays=False, dt_gamma=0, max_steps=1024):
''' march rays to generate points (forward only)
Args:
rays_o/d: float, [N, 3]
bound: float, scalar
density_bitfield: uint8: [CHHH // 8]
C: int
H: int
nears/fars: float, [N]
step_counter: int32, (2), used to count the actual number of generated points.
mean_count: int32, estimated mean steps to accelerate training. (but will randomly drop rays if the actual point count exceeded this threshold.)
perturb: bool
align: int, pad output so its size is dividable by align, set to -1 to disable.
force_all_rays: bool, ignore step_counter and mean_count, always calculate all rays. Useful if rendering the whole image, instead of some rays.
dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
Returns:
xyzs: float, [M, 3], all generated points' coords. (all rays concated, need to use `rays` to extract points belonging to each ray)
dirs: float, [M, 3], all generated points' view dirs.
deltas: float, [M, 2], all generated points' deltas. (first for RGB, second for Depth)
rays: int32, [N, 3], all rays' (index, point_offset, point_count), e.g., xyzs[rays[i, 1]:rays[i, 2]] --> points belonging to rays[i, 0]
'''
if not rays_o.is_cuda: rays_o = rays_o.cuda()
if not rays_d.is_cuda: rays_d = rays_d.cuda()
if not density_bitfield.is_cuda: density_bitfield = density_bitfield.cuda()
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
density_bitfield = density_bitfield.contiguous()
N = rays_o.shape[0] # num rays
M = N * max_steps # init max points number in total
# running average based on previous epoch (mimic `measured_batch_size_before_compaction` in instant-ngp)
# It estimate the max points number to enable faster training, but will lead to random ignored rays if underestimated.
if not force_all_rays and mean_count > 0:
if align > 0:
mean_count += align - mean_count % align
M = mean_count
xyzs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
dirs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
deltas = torch.zeros(M, 2, dtype=rays_o.dtype, device=rays_o.device)
rays = torch.empty(N, 3, dtype=torch.int32, device=rays_o.device) # id, offset, num_steps
if step_counter is None:
step_counter = torch.zeros(2, dtype=torch.int32, device=rays_o.device) # point counter, ray counter
if perturb:
noises = torch.rand(N, dtype=rays_o.dtype, device=rays_o.device)
else:
noises = torch.zeros(N, dtype=rays_o.dtype, device=rays_o.device)
_backend.march_rays_train(rays_o, rays_d, density_bitfield, bound, dt_gamma, max_steps, N, C, H, M, nears, fars, xyzs, dirs, deltas, rays, step_counter, noises) # m is the actually used points number
#print(step_counter, M)
# only used at the first (few) epochs.
if force_all_rays or mean_count <= 0:
m = step_counter[0].item() # D2H copy
if align > 0:
m += align - m % align
xyzs = xyzs[:m]
dirs = dirs[:m]
deltas = deltas[:m]
torch.cuda.empty_cache()
return xyzs, dirs, deltas, rays
march_rays_train = _march_rays_train.apply
class _composite_rays_train(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, sigmas, rgbs, deltas, rays, T_thresh=1e-4):
''' composite rays' rgbs, according to the ray marching formula.
Args:
rgbs: float, [M, 3]
sigmas: float, [M,]
deltas: float, [M, 2]
rays: int32, [N, 3]
Returns:
weights_sum: float, [N,], the alpha channel
depth: float, [N, ], the Depth
image: float, [N, 3], the RGB channel (after multiplying alpha!)
'''
sigmas = sigmas.contiguous()
rgbs = rgbs.contiguous()
M = sigmas.shape[0]
N = rays.shape[0]
weights_sum = torch.empty(N, dtype=sigmas.dtype, device=sigmas.device)
depth = torch.empty(N, dtype=sigmas.dtype, device=sigmas.device)
image = torch.empty(N, 3, dtype=sigmas.dtype, device=sigmas.device)
_backend.composite_rays_train_forward(sigmas, rgbs, deltas, rays, M, N, T_thresh, weights_sum, depth, image)
ctx.save_for_backward(sigmas, rgbs, deltas, rays, weights_sum, depth, image)
ctx.dims = [M, N, T_thresh]
return weights_sum, depth, image
@staticmethod
@custom_bwd
def backward(ctx, grad_weights_sum, grad_depth, grad_image):
# NOTE: grad_depth is not used now! It won't be propagated to sigmas.
grad_weights_sum = grad_weights_sum.contiguous()
grad_image = grad_image.contiguous()
sigmas, rgbs, deltas, rays, weights_sum, depth, image = ctx.saved_tensors
M, N, T_thresh = ctx.dims
grad_sigmas = torch.zeros_like(sigmas)
grad_rgbs = torch.zeros_like(rgbs)
_backend.composite_rays_train_backward(grad_weights_sum, grad_image, sigmas, rgbs, deltas, rays, weights_sum, image, M, N, T_thresh, grad_sigmas, grad_rgbs)
return grad_sigmas, grad_rgbs, None, None, None
composite_rays_train = _composite_rays_train.apply
# ----------------------------------------
# infer functions
# ----------------------------------------
class _march_rays(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, bound, density_bitfield, C, H, near, far, align=-1, perturb=False, dt_gamma=0, max_steps=1024):
''' march rays to generate points (forward only, for inference)
Args:
n_alive: int, number of alive rays
n_step: int, how many steps we march
rays_alive: int, [N], the alive rays' IDs in N (N >= n_alive, but we only use first n_alive)
rays_t: float, [N], the alive rays' time, we only use the first n_alive.
rays_o/d: float, [N, 3]
bound: float, scalar
density_bitfield: uint8: [CHHH // 8]
C: int
H: int
nears/fars: float, [N]
align: int, pad output so its size is dividable by align, set to -1 to disable.
perturb: bool/int, int > 0 is used as the random seed.
dt_gamma: float, called cone_angle in instant-ngp, exponentially accelerate ray marching if > 0. (very significant effect, but generally lead to worse performance)
max_steps: int, max number of sampled points along each ray, also affect min_stepsize.
Returns:
xyzs: float, [n_alive * n_step, 3], all generated points' coords
dirs: float, [n_alive * n_step, 3], all generated points' view dirs.
deltas: float, [n_alive * n_step, 2], all generated points' deltas (here we record two deltas, the first is for RGB, the second for depth).
'''
if not rays_o.is_cuda: rays_o = rays_o.cuda()
if not rays_d.is_cuda: rays_d = rays_d.cuda()
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
M = n_alive * n_step
if align > 0:
M += align - (M % align)
xyzs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
dirs = torch.zeros(M, 3, dtype=rays_o.dtype, device=rays_o.device)
deltas = torch.zeros(M, 2, dtype=rays_o.dtype, device=rays_o.device) # 2 vals, one for rgb, one for depth
if perturb:
# torch.manual_seed(perturb) # test_gui uses spp index as seed
noises = torch.rand(n_alive, dtype=rays_o.dtype, device=rays_o.device)
else:
noises = torch.zeros(n_alive, dtype=rays_o.dtype, device=rays_o.device)
_backend.march_rays(n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, bound, dt_gamma, max_steps, C, H, density_bitfield, near, far, xyzs, dirs, deltas, noises)
return xyzs, dirs, deltas
march_rays = _march_rays.apply
class _composite_rays(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32) # need to cast sigmas & rgbs to float
def forward(ctx, n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, weights_sum, depth, image, T_thresh=1e-2):
''' composite rays' rgbs, according to the ray marching formula. (for inference)
Args:
n_alive: int, number of alive rays
n_step: int, how many steps we march
rays_alive: int, [n_alive], the alive rays' IDs in N (N >= n_alive)
rays_t: float, [N], the alive rays' time
sigmas: float, [n_alive * n_step,]
rgbs: float, [n_alive * n_step, 3]
deltas: float, [n_alive * n_step, 2], all generated points' deltas (here we record two deltas, the first is for RGB, the second for depth).
In-place Outputs:
weights_sum: float, [N,], the alpha channel
depth: float, [N,], the depth value
image: float, [N, 3], the RGB channel (after multiplying alpha!)
'''
_backend.composite_rays(n_alive, n_step, T_thresh, rays_alive, rays_t, sigmas, rgbs, deltas, weights_sum, depth, image)
return tuple()
composite_rays = _composite_rays.apply
================================================
FILE: 3drecon/raymarching/setup.py
================================================
import os
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
_src_path = os.path.dirname(os.path.abspath(__file__))
nvcc_flags = [
'-O3', '-std=c++14',
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__',
]
if os.name == "posix":
c_flags = ['-O3', '-std=c++14']
elif os.name == "nt":
c_flags = ['/O2', '/std:c++17']
# find cl.exe
def find_cl_path():
import glob
for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
paths = sorted(glob.glob(r"C:\\Program Files (x86)\\Microsoft Visual Studio\\*\\%s\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64" % edition), reverse=True)
if paths:
return paths[0]
# If cl.exe is not on path, try to find it.
if os.system("where cl.exe >nul 2>nul") != 0:
cl_path = find_cl_path()
if cl_path is None:
raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
os.environ["PATH"] += ";" + cl_path
'''
Usage:
python setup.py build_ext --inplace # build extensions locally, do not install (only can be used from the parent directory)
python setup.py install # build extensions and install (copy) to PATH.
pip install . # ditto but better (e.g., dependency & metadata handling)
python setup.py develop # build extensions and install (symbolic) to PATH.
pip install -e . # ditto but better (e.g., dependency & metadata handling)
'''
setup(
name='raymarching', # package name, import this to use python API
ext_modules=[
CUDAExtension(
name='_raymarching', # extension name, import this to use CUDA API
sources=[os.path.join(_src_path, 'src', f) for f in [
'raymarching.cu',
'bindings.cpp',
]],
extra_compile_args={
'cxx': c_flags,
'nvcc': nvcc_flags,
}
),
],
cmdclass={
'build_ext': BuildExtension,
}
)
================================================
FILE: 3drecon/raymarching/src/bindings.cpp
================================================
#include <torch/extension.h>
#include "raymarching.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// utils
m.def("packbits", &packbits, "packbits (CUDA)");
m.def("near_far_from_aabb", &near_far_from_aabb, "near_far_from_aabb (CUDA)");
m.def("sph_from_ray", &sph_from_ray, "sph_from_ray (CUDA)");
m.def("morton3D", &morton3D, "morton3D (CUDA)");
m.def("morton3D_invert", &morton3D_invert, "morton3D_invert (CUDA)");
// train
m.def("march_rays_train", &march_rays_train, "march_rays_train (CUDA)");
m.def("composite_rays_train_forward", &composite_rays_train_forward, "composite_rays_train_forward (CUDA)");
m.def("composite_rays_train_backward", &composite_rays_train_backward, "composite_rays_train_backward (CUDA)");
// infer
m.def("march_rays", &march_rays, "march rays (CUDA)");
m.def("composite_rays", &composite_rays, "composite rays (CUDA)");
}
================================================
FILE: 3drecon/raymarching/src/raymarching.cu
================================================
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/torch.h>
#include <cstdio>
#include <stdint.h>
#include <stdexcept>
#include <limits>
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor")
#define CHECK_IS_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor")
#define CHECK_IS_FLOATING(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float || x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::Double, #x " must be a floating tensor")
inline constexpr __device__ float SQRT3() { return 1.7320508075688772f; }
inline constexpr __device__ float RSQRT3() { return 0.5773502691896258f; }
inline constexpr __device__ float PI() { return 3.141592653589793f; }
inline constexpr __device__ float RPI() { return 0.3183098861837907f; }
template <typename T>
inline __host__ __device__ T div_round_up(T val, T divisor) {
return (val + divisor - 1) / divisor;
}
inline __host__ __device__ float signf(const float x) {
return copysignf(1.0, x);
}
inline __host__ __device__ float clamp(const float x, const float min, const float max) {
return fminf(max, fmaxf(min, x));
}
inline __host__ __device__ void swapf(float& a, float& b) {
float c = a; a = b; b = c;
}
inline __device__ int mip_from_pos(const float x, const float y, const float z, const float max_cascade) {
const float mx = fmaxf(fabsf(x), fmaxf(fabs(y), fabs(z)));
int exponent;
frexpf(mx, &exponent); // [0, 0.5) --> -1, [0.5, 1) --> 0, [1, 2) --> 1, [2, 4) --> 2, ...
return fminf(max_cascade - 1, fmaxf(0, exponent));
}
inline __device__ int mip_from_dt(const float dt, const float H, const float max_cascade) {
const float mx = dt * H * 0.5;
int exponent;
frexpf(mx, &exponent);
return fminf(max_cascade - 1, fmaxf(0, exponent));
}
inline __host__ __device__ uint32_t __expand_bits(uint32_t v)
{
v = (v * 0x00010001u) & 0xFF0000FFu;
v = (v * 0x00000101u) & 0x0F00F00Fu;
v = (v * 0x00000011u) & 0xC30C30C3u;
v = (v * 0x00000005u) & 0x49249249u;
return v;
}
inline __host__ __device__ uint32_t __morton3D(uint32_t x, uint32_t y, uint32_t z)
{
uint32_t xx = __expand_bits(x);
uint32_t yy = __expand_bits(y);
uint32_t zz = __expand_bits(z);
return xx | (yy << 1) | (zz << 2);
}
inline __host__ __device__ uint32_t __morton3D_invert(uint32_t x)
{
x = x & 0x49249249;
x = (x | (x >> 2)) & 0xc30c30c3;
x = (x | (x >> 4)) & 0x0f00f00f;
x = (x | (x >> 8)) & 0xff0000ff;
x = (x | (x >> 16)) & 0x0000ffff;
return x;
}
////////////////////////////////////////////////////
///////////// utils /////////////
////////////////////////////////////////////////////
// rays_o/d: [N, 3]
// nears/fars: [N]
// scalar_t should always be float in use.
template <typename scalar_t>
__global__ void kernel_near_far_from_aabb(
const scalar_t * __restrict__ rays_o,
const scalar_t * __restrict__ rays_d,
const scalar_t * __restrict__ aabb,
const uint32_t N,
const float min_near,
scalar_t * nears, scalar_t * fars
) {
// parallel per ray
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
rays_o += n * 3;
rays_d += n * 3;
const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;
// get near far (assume cube scene)
float near = (aabb[0] - ox) * rdx;
float far = (aabb[3] - ox) * rdx;
if (near > far) swapf(near, far);
float near_y = (aabb[1] - oy) * rdy;
float far_y = (aabb[4] - oy) * rdy;
if (near_y > far_y) swapf(near_y, far_y);
if (near > far_y || near_y > far) {
nears[n] = fars[n] = std::numeric_limits<scalar_t>::max();
return;
}
if (near_y > near) near = near_y;
if (far_y < far) far = far_y;
float near_z = (aabb[2] - oz) * rdz;
float far_z = (aabb[5] - oz) * rdz;
if (near_z > far_z) swapf(near_z, far_z);
if (near > far_z || near_z > far) {
nears[n] = fars[n] = std::numeric_limits<scalar_t>::max();
return;
}
if (near_z > near) near = near_z;
if (far_z < far) far = far_z;
if (near < min_near) near = min_near;
nears[n] = near;
fars[n] = far;
}
void near_far_from_aabb(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor aabb, const uint32_t N, const float min_near, at::Tensor nears, at::Tensor fars) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
rays_o.scalar_type(), "near_far_from_aabb", ([&] {
kernel_near_far_from_aabb<<<div_round_up(N, N_THREAD), N_THREAD>>>(rays_o.data_ptr<scalar_t>(), rays_d.data_ptr<scalar_t>(), aabb.data_ptr<scalar_t>(), N, min_near, nears.data_ptr<scalar_t>(), fars.data_ptr<scalar_t>());
}));
}
// rays_o/d: [N, 3]
// radius: float
// coords: [N, 2]
template <typename scalar_t>
__global__ void kernel_sph_from_ray(
const scalar_t * __restrict__ rays_o,
const scalar_t * __restrict__ rays_d,
const float radius,
const uint32_t N,
scalar_t * coords
) {
// parallel per ray
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
rays_o += n * 3;
rays_d += n * 3;
coords += n * 2;
const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;
// solve t from || o + td || = radius
const float A = dx * dx + dy * dy + dz * dz;
const float B = ox * dx + oy * dy + oz * dz; // in fact B / 2
const float C = ox * ox + oy * oy + oz * oz - radius * radius;
const float t = (- B + sqrtf(B * B - A * C)) / A; // always use the larger solution (positive)
// solve theta, phi (assume y is the up axis)
const float x = ox + t * dx, y = oy + t * dy, z = oz + t * dz;
const float theta = atan2(sqrtf(x * x + z * z), y); // [0, PI)
const float phi = atan2(z, x); // [-PI, PI)
// normalize to [-1, 1]
coords[0] = 2 * theta * RPI() - 1;
coords[1] = phi * RPI();
}
void sph_from_ray(const at::Tensor rays_o, const at::Tensor rays_d, const float radius, const uint32_t N, at::Tensor coords) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
rays_o.scalar_type(), "sph_from_ray", ([&] {
kernel_sph_from_ray<<<div_round_up(N, N_THREAD), N_THREAD>>>(rays_o.data_ptr<scalar_t>(), rays_d.data_ptr<scalar_t>(), radius, N, coords.data_ptr<scalar_t>());
}));
}
// coords: int32, [N, 3]
// indices: int32, [N]
__global__ void kernel_morton3D(
const int * __restrict__ coords,
const uint32_t N,
int * indices
) {
// parallel
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
coords += n * 3;
indices[n] = __morton3D(coords[0], coords[1], coords[2]);
}
void morton3D(const at::Tensor coords, const uint32_t N, at::Tensor indices) {
static constexpr uint32_t N_THREAD = 128;
kernel_morton3D<<<div_round_up(N, N_THREAD), N_THREAD>>>(coords.data_ptr<int>(), N, indices.data_ptr<int>());
}
// indices: int32, [N]
// coords: int32, [N, 3]
__global__ void kernel_morton3D_invert(
const int * __restrict__ indices,
const uint32_t N,
int * coords
) {
// parallel
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
coords += n * 3;
const int ind = indices[n];
coords[0] = __morton3D_invert(ind >> 0);
coords[1] = __morton3D_invert(ind >> 1);
coords[2] = __morton3D_invert(ind >> 2);
}
void morton3D_invert(const at::Tensor indices, const uint32_t N, at::Tensor coords) {
static constexpr uint32_t N_THREAD = 128;
kernel_morton3D_invert<<<div_round_up(N, N_THREAD), N_THREAD>>>(indices.data_ptr<int>(), N, coords.data_ptr<int>());
}
// grid: float, [C, H, H, H]
// N: int, C * H * H * H / 8
// density_thresh: float
// bitfield: uint8, [N]
template <typename scalar_t>
__global__ void kernel_packbits(
const scalar_t * __restrict__ grid,
const uint32_t N,
const float density_thresh,
uint8_t * bitfield
) {
// parallel per byte
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
grid += n * 8;
uint8_t bits = 0;
#pragma unroll
for (uint8_t i = 0; i < 8; i++) {
bits |= (grid[i] > density_thresh) ? ((uint8_t)1 << i) : 0;
}
bitfield[n] = bits;
}
void packbits(const at::Tensor grid, const uint32_t N, const float density_thresh, at::Tensor bitfield) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
grid.scalar_type(), "packbits", ([&] {
kernel_packbits<<<div_round_up(N, N_THREAD), N_THREAD>>>(grid.data_ptr<scalar_t>(), N, density_thresh, bitfield.data_ptr<uint8_t>());
}));
}
////////////////////////////////////////////////////
///////////// training /////////////
////////////////////////////////////////////////////
// rays_o/d: [N, 3]
// grid: [CHHH / 8]
// xyzs, dirs, deltas: [M, 3], [M, 3], [M, 2]
// dirs: [M, 3]
// rays: [N, 3], idx, offset, num_steps
template <typename scalar_t>
__global__ void kernel_march_rays_train(
const scalar_t * __restrict__ rays_o,
const scalar_t * __restrict__ rays_d,
const uint8_t * __restrict__ grid,
const float bound,
const float dt_gamma, const uint32_t max_steps,
const uint32_t N, const uint32_t C, const uint32_t H, const uint32_t M,
const scalar_t* __restrict__ nears,
const scalar_t* __restrict__ fars,
scalar_t * xyzs, scalar_t * dirs, scalar_t * deltas,
int * rays,
int * counter,
const scalar_t* __restrict__ noises
) {
// parallel per ray
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
rays_o += n * 3;
rays_d += n * 3;
// ray marching
const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;
const float rH = 1 / (float)H;
const float H3 = H * H * H;
const float near = nears[n];
const float far = fars[n];
const float noise = noises[n];
const float dt_min = 2 * SQRT3() / max_steps;
const float dt_max = 2 * SQRT3() * (1 << (C - 1)) / H;
float t0 = near;
// perturb
t0 += clamp(t0 * dt_gamma, dt_min, dt_max) * noise;
// first pass: estimation of num_steps
float t = t0;
uint32_t num_steps = 0;
//if (t < far) printf("valid ray %d t=%f near=%f far=%f \n", n, t, near, far);
while (t < far && num_steps < max_steps) {
// current point
const float x = clamp(ox + t * dx, -bound, bound);
const float y = clamp(oy + t * dy, -bound, bound);
const float z = clamp(oz + t * dz, -bound, bound);
const float dt = clamp(t * dt_gamma, dt_min, dt_max);
// get mip level
const int level = max(mip_from_pos(x, y, z, C), mip_from_dt(dt, H, C)); // range in [0, C - 1]
const float mip_bound = fminf(scalbnf(1.0f, level), bound);
const float mip_rbound = 1 / mip_bound;
// convert to nearest grid position
const int nx = clamp(0.5 * (x * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int ny = clamp(0.5 * (y * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int nz = clamp(0.5 * (z * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const uint32_t index = level * H3 + __morton3D(nx, ny, nz);
const bool occ = grid[index / 8] & (1 << (index % 8));
// if occpuied, advance a small step, and write to output
//if (n == 0) printf("t=%f density=%f vs thresh=%f step=%d\n", t, density, density_thresh, num_steps);
if (occ) {
num_steps++;
t += dt;
// else, skip a large step (basically skip a voxel grid)
} else {
// calc distance to next voxel
const float tx = (((nx + 0.5f + 0.5f * signf(dx)) * rH * 2 - 1) * mip_bound - x) * rdx;
const float ty = (((ny + 0.5f + 0.5f * signf(dy)) * rH * 2 - 1) * mip_bound - y) * rdy;
const float tz = (((nz + 0.5f + 0.5f * signf(dz)) * rH * 2 - 1) * mip_bound - z) * rdz;
const float tt = t + fmaxf(0.0f, fminf(tx, fminf(ty, tz)));
// step until next voxel
do {
t += clamp(t * dt_gamma, dt_min, dt_max);
} while (t < tt);
}
}
//printf("[n=%d] num_steps=%d, near=%f, far=%f, dt=%f, max_steps=%f\n", n, num_steps, near, far, dt_min, (far - near) / dt_min);
// second pass: really locate and write points & dirs
uint32_t point_index = atomicAdd(counter, num_steps);
uint32_t ray_index = atomicAdd(counter + 1, 1);
//printf("[n=%d] num_steps=%d, point_index=%d, ray_index=%d\n", n, num_steps, point_index, ray_index);
// write rays
rays[ray_index * 3] = n;
rays[ray_index * 3 + 1] = point_index;
rays[ray_index * 3 + 2] = num_steps;
if (num_steps == 0) return;
if (point_index + num_steps > M) return;
xyzs += point_index * 3;
dirs += point_index * 3;
deltas += point_index * 2;
t = t0;
uint32_t step = 0;
float last_t = t;
while (t < far && step < num_steps) {
// current point
const float x = clamp(ox + t * dx, -bound, bound);
const float y = clamp(oy + t * dy, -bound, bound);
const float z = clamp(oz + t * dz, -bound, bound);
const float dt = clamp(t * dt_gamma, dt_min, dt_max);
// get mip level
const int level = max(mip_from_pos(x, y, z, C), mip_from_dt(dt, H, C)); // range in [0, C - 1]
const float mip_bound = fminf(scalbnf(1.0f, level), bound);
const float mip_rbound = 1 / mip_bound;
// convert to nearest grid position
const int nx = clamp(0.5 * (x * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int ny = clamp(0.5 * (y * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int nz = clamp(0.5 * (z * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
// query grid
const uint32_t index = level * H3 + __morton3D(nx, ny, nz);
const bool occ = grid[index / 8] & (1 << (index % 8));
// if occpuied, advance a small step, and write to output
if (occ) {
// write step
xyzs[0] = x;
xyzs[1] = y;
xyzs[2] = z;
dirs[0] = dx;
dirs[1] = dy;
dirs[2] = dz;
t += dt;
deltas[0] = dt;
deltas[1] = t - last_t; // used to calc depth
last_t = t;
xyzs += 3;
dirs += 3;
deltas += 2;
step++;
// else, skip a large step (basically skip a voxel grid)
} else {
// calc distance to next voxel
const float tx = (((nx + 0.5f + 0.5f * signf(dx)) * rH * 2 - 1) * mip_bound - x) * rdx;
const float ty = (((ny + 0.5f + 0.5f * signf(dy)) * rH * 2 - 1) * mip_bound - y) * rdy;
const float tz = (((nz + 0.5f + 0.5f * signf(dz)) * rH * 2 - 1) * mip_bound - z) * rdz;
const float tt = t + fmaxf(0.0f, fminf(tx, fminf(ty, tz)));
// step until next voxel
do {
t += clamp(t * dt_gamma, dt_min, dt_max);
} while (t < tt);
}
}
}
void march_rays_train(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor grid, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t N, const uint32_t C, const uint32_t H, const uint32_t M, const at::Tensor nears, const at::Tensor fars, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor rays, at::Tensor counter, at::Tensor noises) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
rays_o.scalar_type(), "march_rays_train", ([&] {
kernel_march_rays_train<<<div_round_up(N, N_THREAD), N_THREAD>>>(rays_o.data_ptr<scalar_t>(), rays_d.data_ptr<scalar_t>(), grid.data_ptr<uint8_t>(), bound, dt_gamma, max_steps, N, C, H, M, nears.data_ptr<scalar_t>(), fars.data_ptr<scalar_t>(), xyzs.data_ptr<scalar_t>(), dirs.data_ptr<scalar_t>(), deltas.data_ptr<scalar_t>(), rays.data_ptr<int>(), counter.data_ptr<int>(), noises.data_ptr<scalar_t>());
}));
}
// sigmas: [M]
// rgbs: [M, 3]
// deltas: [M, 2]
// rays: [N, 3], idx, offset, num_steps
// weights_sum: [N], final pixel alpha
// depth: [N,]
// image: [N, 3]
template <typename scalar_t>
__global__ void kernel_composite_rays_train_forward(
const scalar_t * __restrict__ sigmas,
const scalar_t * __restrict__ rgbs,
const scalar_t * __restrict__ deltas,
const int * __restrict__ rays,
const uint32_t M, const uint32_t N, const float T_thresh,
scalar_t * weights_sum,
scalar_t * depth,
scalar_t * image
) {
// parallel per ray
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
uint32_t index = rays[n * 3];
uint32_t offset = rays[n * 3 + 1];
uint32_t num_steps = rays[n * 3 + 2];
// empty ray, or ray that exceed max step count.
if (num_steps == 0 || offset + num_steps > M) {
weights_sum[index] = 0;
depth[index] = 0;
image[index * 3] = 0;
image[index * 3 + 1] = 0;
image[index * 3 + 2] = 0;
return;
}
sigmas += offset;
rgbs += offset * 3;
deltas += offset * 2;
// accumulate
uint32_t step = 0;
scalar_t T = 1.0f;
scalar_t r = 0, g = 0, b = 0, ws = 0, t = 0, d = 0;
while (step < num_steps) {
const scalar_t alpha = 1.0f - __expf(- sigmas[0] * deltas[0]);
const scalar_t weight = alpha * T;
r += weight * rgbs[0];
g += weight * rgbs[1];
b += weight * rgbs[2];
t += deltas[1]; // real delta
d += weight * t;
ws += weight;
T *= 1.0f - alpha;
// minimal remained transmittence
if (T < T_thresh) break;
//printf("[n=%d] num_steps=%d, alpha=%f, w=%f, T=%f, sum_dt=%f, d=%f\n", n, step, alpha, weight, T, sum_delta, d);
// locate
sigmas++;
rgbs += 3;
deltas += 2;
step++;
}
//printf("[n=%d] rgb=(%f, %f, %f), d=%f\n", n, r, g, b, d);
// write
weights_sum[index] = ws; // weights_sum
depth[index] = d;
image[index * 3] = r;
image[index * 3 + 1] = g;
image[index * 3 + 2] = b;
}
void composite_rays_train_forward(const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor weights_sum, at::Tensor depth, at::Tensor image) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
sigmas.scalar_type(), "composite_rays_train_forward", ([&] {
kernel_composite_rays_train_forward<<<div_round_up(N, N_THREAD), N_THREAD>>>(sigmas.data_ptr<scalar_t>(), rgbs.data_ptr<scalar_t>(), deltas.data_ptr<scalar_t>(), rays.data_ptr<int>(), M, N, T_thresh, weights_sum.data_ptr<scalar_t>(), depth.data_ptr<scalar_t>(), image.data_ptr<scalar_t>());
}));
}
// grad_weights_sum: [N,]
// grad: [N, 3]
// sigmas: [M]
// rgbs: [M, 3]
// deltas: [M, 2]
// rays: [N, 3], idx, offset, num_steps
// weights_sum: [N,], weights_sum here
// image: [N, 3]
// grad_sigmas: [M]
// grad_rgbs: [M, 3]
template <typename scalar_t>
__global__ void kernel_composite_rays_train_backward(
const scalar_t * __restrict__ grad_weights_sum,
const scalar_t * __restrict__ grad_image,
const scalar_t * __restrict__ sigmas,
const scalar_t * __restrict__ rgbs,
const scalar_t * __restrict__ deltas,
const int * __restrict__ rays,
const scalar_t * __restrict__ weights_sum,
const scalar_t * __restrict__ image,
const uint32_t M, const uint32_t N, const float T_thresh,
scalar_t * grad_sigmas,
scalar_t * grad_rgbs
) {
// parallel per ray
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= N) return;
// locate
uint32_t index = rays[n * 3];
uint32_t offset = rays[n * 3 + 1];
uint32_t num_steps = rays[n * 3 + 2];
if (num_steps == 0 || offset + num_steps > M) return;
grad_weights_sum += index;
grad_image += index * 3;
weights_sum += index;
image += index * 3;
sigmas += offset;
rgbs += offset * 3;
deltas += offset * 2;
grad_sigmas += offset;
grad_rgbs += offset * 3;
// accumulate
uint32_t step = 0;
scalar_t T = 1.0f;
const scalar_t r_final = image[0], g_final = image[1], b_final = image[2], ws_final = weights_sum[0];
scalar_t r = 0, g = 0, b = 0, ws = 0;
while (step < num_steps) {
const scalar_t alpha = 1.0f - __expf(- sigmas[0] * deltas[0]);
const scalar_t weight = alpha * T;
r += weight * rgbs[0];
g += weight * rgbs[1];
b += weight * rgbs[2];
ws += weight;
T *= 1.0f - alpha;
// check https://note.kiui.moe/others/nerf_gradient/ for the gradient calculation.
// write grad_rgbs
grad_rgbs[0] = grad_image[0] * weight;
grad_rgbs[1] = grad_image[1] * weight;
grad_rgbs[2] = grad_image[2] * weight;
// write grad_sigmas
grad_sigmas[0] = deltas[0] * (
grad_image[0] * (T * rgbs[0] - (r_final - r)) +
grad_image[1] * (T * rgbs[1] - (g_final - g)) +
grad_image[2] * (T * rgbs[2] - (b_final - b)) +
grad_weights_sum[0] * (1 - ws_final)
);
//printf("[n=%d] num_steps=%d, T=%f, grad_sigmas=%f, r_final=%f, r=%f\n", n, step, T, grad_sigmas[0], r_final, r);
// minimal remained transmittence
if (T < T_thresh) break;
// locate
sigmas++;
rgbs += 3;
deltas += 2;
grad_sigmas++;
grad_rgbs += 3;
step++;
}
}
void composite_rays_train_backward(const at::Tensor grad_weights_sum, const at::Tensor grad_image, const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const at::Tensor weights_sum, const at::Tensor image, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor grad_sigmas, at::Tensor grad_rgbs) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
grad_image.scalar_type(), "composite_rays_train_backward", ([&] {
kernel_composite_rays_train_backward<<<div_round_up(N, N_THREAD), N_THREAD>>>(grad_weights_sum.data_ptr<scalar_t>(), grad_image.data_ptr<scalar_t>(), sigmas.data_ptr<scalar_t>(), rgbs.data_ptr<scalar_t>(), deltas.data_ptr<scalar_t>(), rays.data_ptr<int>(), weights_sum.data_ptr<scalar_t>(), image.data_ptr<scalar_t>(), M, N, T_thresh, grad_sigmas.data_ptr<scalar_t>(), grad_rgbs.data_ptr<scalar_t>());
}));
}
////////////////////////////////////////////////////
///////////// infernce /////////////
////////////////////////////////////////////////////
template <typename scalar_t>
__global__ void kernel_march_rays(
const uint32_t n_alive,
const uint32_t n_step,
const int* __restrict__ rays_alive,
const scalar_t* __restrict__ rays_t,
const scalar_t* __restrict__ rays_o,
const scalar_t* __restrict__ rays_d,
const float bound,
const float dt_gamma, const uint32_t max_steps,
const uint32_t C, const uint32_t H,
const uint8_t * __restrict__ grid,
const scalar_t* __restrict__ nears,
const scalar_t* __restrict__ fars,
scalar_t* xyzs, scalar_t* dirs, scalar_t* deltas,
const scalar_t* __restrict__ noises
) {
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= n_alive) return;
const int index = rays_alive[n]; // ray id
const float noise = noises[n];
// locate
rays_o += index * 3;
rays_d += index * 3;
xyzs += n * n_step * 3;
dirs += n * n_step * 3;
deltas += n * n_step * 2;
const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2];
const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2];
const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz;
const float rH = 1 / (float)H;
const float H3 = H * H * H;
float t = rays_t[index]; // current ray's t
const float near = nears[index], far = fars[index];
const float dt_min = 2 * SQRT3() / max_steps;
const float dt_max = 2 * SQRT3() * (1 << (C - 1)) / H;
// march for n_step steps, record points
uint32_t step = 0;
// introduce some randomness
t += clamp(t * dt_gamma, dt_min, dt_max) * noise;
float last_t = t;
while (t < far && step < n_step) {
// current point
const float x = clamp(ox + t * dx, -bound, bound);
const float y = clamp(oy + t * dy, -bound, bound);
const float z = clamp(oz + t * dz, -bound, bound);
const float dt = clamp(t * dt_gamma, dt_min, dt_max);
// get mip level
const int level = max(mip_from_pos(x, y, z, C), mip_from_dt(dt, H, C)); // range in [0, C - 1]
const float mip_bound = fminf(scalbnf(1, level), bound);
const float mip_rbound = 1 / mip_bound;
// convert to nearest grid position
const int nx = clamp(0.5 * (x * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int ny = clamp(0.5 * (y * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const int nz = clamp(0.5 * (z * mip_rbound + 1) * H, 0.0f, (float)(H - 1));
const uint32_t index = level * H3 + __morton3D(nx, ny, nz);
const bool occ = grid[index / 8] & (1 << (index % 8));
// if occpuied, advance a small step, and write to output
if (occ) {
// write step
xyzs[0] = x;
xyzs[1] = y;
xyzs[2] = z;
dirs[0] = dx;
dirs[1] = dy;
dirs[2] = dz;
// calc dt
t += dt;
deltas[0] = dt;
deltas[1] = t - last_t; // used to calc depth
last_t = t;
// step
xyzs += 3;
dirs += 3;
deltas += 2;
step++;
// else, skip a large step (basically skip a voxel grid)
} else {
// calc distance to next voxel
const float tx = (((nx + 0.5f + 0.5f * signf(dx)) * rH * 2 - 1) * mip_bound - x) * rdx;
const float ty = (((ny + 0.5f + 0.5f * signf(dy)) * rH * 2 - 1) * mip_bound - y) * rdy;
const float tz = (((nz + 0.5f + 0.5f * signf(dz)) * rH * 2 - 1) * mip_bound - z) * rdz;
const float tt = t + fmaxf(0.0f, fminf(tx, fminf(ty, tz)));
// step until next voxel
do {
t += clamp(t * dt_gamma, dt_min, dt_max);
} while (t < tt);
}
}
}
void march_rays(const uint32_t n_alive, const uint32_t n_step, const at::Tensor rays_alive, const at::Tensor rays_t, const at::Tensor rays_o, const at::Tensor rays_d, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t C, const uint32_t H, const at::Tensor grid, const at::Tensor near, const at::Tensor far, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor noises) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
rays_o.scalar_type(), "march_rays", ([&] {
kernel_march_rays<<<div_round_up(n_alive, N_THREAD), N_THREAD>>>(n_alive, n_step, rays_alive.data_ptr<int>(), rays_t.data_ptr<scalar_t>(), rays_o.data_ptr<scalar_t>(), rays_d.data_ptr<scalar_t>(), bound, dt_gamma, max_steps, C, H, grid.data_ptr<uint8_t>(), near.data_ptr<scalar_t>(), far.data_ptr<scalar_t>(), xyzs.data_ptr<scalar_t>(), dirs.data_ptr<scalar_t>(), deltas.data_ptr<scalar_t>(), noises.data_ptr<scalar_t>());
}));
}
template <typename scalar_t>
__global__ void kernel_composite_rays(
const uint32_t n_alive,
const uint32_t n_step,
const float T_thresh,
int* rays_alive,
scalar_t* rays_t,
const scalar_t* __restrict__ sigmas,
const scalar_t* __restrict__ rgbs,
const scalar_t* __restrict__ deltas,
scalar_t* weights_sum, scalar_t* depth, scalar_t* image
) {
const uint32_t n = threadIdx.x + blockIdx.x * blockDim.x;
if (n >= n_alive) return;
const int index = rays_alive[n]; // ray id
// locate
sigmas += n * n_step;
rgbs += n * n_step * 3;
deltas += n * n_step * 2;
rays_t += index;
weights_sum += index;
depth += index;
image += index * 3;
scalar_t t = rays_t[0]; // current ray's t
scalar_t weight_sum = weights_sum[0];
scalar_t d = depth[0];
scalar_t r = image[0];
scalar_t g = image[1];
scalar_t b = image[2];
// accumulate
uint32_t step = 0;
while (step < n_step) {
// ray is terminated if delta == 0
if (deltas[0] == 0) break;
const scalar_t alpha = 1.0f - __expf(- sigmas[0] * deltas[0]);
/*
T_0 = 1; T_i = \prod_{j=0}^{i-1} (1 - alpha_j)
w_i = alpha_i * T_i
-->
T_i = 1 - \sum_{j=0}^{i-1} w_j
*/
const scalar_t T = 1 - weight_sum;
const scalar_t weight = alpha * T;
weight_sum += weight;
t += deltas[1]; // real delta
d += weight * t;
r += weight * rgbs[0];
g += weight * rgbs[1];
b += weight * rgbs[2];
//printf("[n=%d] num_steps=%d, alpha=%f, w=%f, T=%f, sum_dt=%f, d=%f\n", n, step, alpha, weight, T, sum_delta, d);
// ray is terminated if T is too small
// use a larger bound to further accelerate inference
if (T < T_thresh) break;
// locate
sigmas++;
rgbs += 3;
deltas += 2;
step++;
}
//printf("[n=%d] rgb=(%f, %f, %f), d=%f\n", n, r, g, b, d);
// rays_alive = -1 means ray is terminated early.
if (step < n_step) {
rays_alive[n] = -1;
} else {
rays_t[0] = t;
}
weights_sum[0] = weight_sum; // this is the thing I needed!
depth[0] = d;
image[0] = r;
image[1] = g;
image[2] = b;
}
void composite_rays(const uint32_t n_alive, const uint32_t n_step, const float T_thresh, at::Tensor rays_alive, at::Tensor rays_t, at::Tensor sigmas, at::Tensor rgbs, at::Tensor deltas, at::Tensor weights, at::Tensor depth, at::Tensor image) {
static constexpr uint32_t N_THREAD = 128;
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
image.scalar_type(), "composite_rays", ([&] {
kernel_composite_rays<<<div_round_up(n_alive, N_THREAD), N_THREAD>>>(n_alive, n_step, T_thresh, rays_alive.data_ptr<int>(), rays_t.data_ptr<scalar_t>(), sigmas.data_ptr<scalar_t>(), rgbs.data_ptr<scalar_t>(), deltas.data_ptr<scalar_t>(), weights.data_ptr<scalar_t>(), depth.data_ptr<scalar_t>(), image.data_ptr<scalar_t>());
}));
}
================================================
FILE: 3drecon/raymarching/src/raymarching.h
================================================
#pragma once
#include <stdint.h>
#include <torch/torch.h>
void near_far_from_aabb(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor aabb, const uint32_t N, const float min_near, at::Tensor nears, at::Tensor fars);
void sph_from_ray(const at::Tensor rays_o, const at::Tensor rays_d, const float radius, const uint32_t N, at::Tensor coords);
void morton3D(const at::Tensor coords, const uint32_t N, at::Tensor indices);
void morton3D_invert(const at::Tensor indices, const uint32_t N, at::Tensor coords);
void packbits(const at::Tensor grid, const uint32_t N, const float density_thresh, at::Tensor bitfield);
void march_rays_train(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor grid, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t N, const uint32_t C, const uint32_t H, const uint32_t M, const at::Tensor nears, const at::Tensor fars, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor rays, at::Tensor counter, at::Tensor noises);
void composite_rays_train_forward(const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor weights_sum, at::Tensor depth, at::Tensor image);
void composite_rays_train_backward(const at::Tensor grad_weights_sum, const at::Tensor grad_image, const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const at::Tensor weights_sum, const at::Tensor image, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor grad_sigmas, at::Tensor grad_rgbs);
void march_rays(const uint32_t n_alive, const uint32_t n_step, const at::Tensor rays_alive, const at::Tensor rays_t, const at::Tensor rays_o, const at::Tensor rays_d, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t C, const uint32_t H, const at::Tensor grid, const at::Tensor nears, const at::Tensor fars, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor noises);
void composite_rays(const uint32_t n_alive, const uint32_t n_step, const float T_thresh, at::Tensor rays_alive, at::Tensor rays_t, at::Tensor sigmas, at::Tensor rgbs, at::Tensor deltas, at::Tensor weights_sum, at::Tensor depth, at::Tensor image);
================================================
FILE: 3drecon/renderer/agg_net.py
================================================
import torch.nn.functional as F
import torch.nn as nn
import torch
def weights_init(m):
if isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight.data)
if m.bias is not None:
nn.init.zeros_(m.bias.data)
class NeRF(nn.Module):
def __init__(self, vol_n=8+8, feat_ch=8+16+32+3, hid_n=64):
super(NeRF, self).__init__()
self.hid_n = hid_n
self.agg = Agg(feat_ch)
self.lr0 = nn.Sequential(nn.Linear(vol_n+16, hid_n), nn.ReLU())
self.sigma = nn.Sequential(nn.Linear(hid_n, 1), nn.Softplus())
self.color = nn.Sequential(
nn.Linear(16+vol_n+feat_ch+hid_n+4, hid_n), # agg_feats+vox_feat+img_feat+lr0_feats+dir
nn.ReLU(),
nn.Linear(hid_n, 1)
)
self.lr0.apply(weights_init)
self.sigma.apply(weights_init)
self.color.apply(weights_init)
def forward(self, vox_feat, img_feat_rgb_dir, source_img_mask):
# assert torch.sum(torch.sum(source_img_mask,1)<2)==0
b, d, n, _ = img_feat_rgb_dir.shape # b,d,n,f=8+16+32+3+4
agg_feat = self.agg(img_feat_rgb_dir, source_img_mask) # b,d,f=16
x = self.lr0(torch.cat((vox_feat, agg_feat), dim=-1)) # b,d,f=64
sigma = self.sigma(x) # b,d,1
x = torch.cat((x, vox_feat, agg_feat), dim=-1) # b,d,f=16+16+64
x = x.view(b, d, 1, x.shape[-1]).repeat(1, 1, n, 1)
x = torch.cat((x, img_feat_rgb_dir), dim=-1)
logits = self.color(x)
source_img_mask_ = source_img_mask.reshape(b, 1, n, 1).repeat(1, logits.shape[1], 1, 1) == 0
logits[source_img_mask_] = -1e7
color_weight = F.softmax(logits, dim=-2)
color = torch.sum((img_feat_rgb_dir[..., -7:-4] * color_weight), dim=-2)
return color, sigma
class Agg(nn.Module):
def __init__(self, feat_ch):
super(Agg, self).__init__()
self.feat_ch = feat_ch
self.view_fc = nn.Sequential(nn.Linear(4, feat_ch), nn.ReLU())
self.view_fc.apply(weights_init)
self.global_fc = nn.Sequential(nn.Linear(feat_ch*3, 32), nn.ReLU())
self.agg_w_fc = nn.Linear(32, 1)
self.fc = nn.Linear(32, 16)
self.global_fc.apply(weights_init)
self.agg_w_fc.apply(weights_init)
self.fc.apply(weights_init)
def masked_mean_var(self, img_feat_rgb, source_img_mask):
# img_feat_rgb: b,d,n,f source_img_mask: b,n
b, n = source_img_mask.shape
source_img_mask = source_img_mask.view(b, 1, n, 1)
mean = torch.sum(source_img_mask * img_feat_rgb, dim=-2)/ (torch.sum(source_img_mask, dim=-2) + 1e-5)
var = torch.sum((img_feat_rgb - mean.unsqueeze(-2)) ** 2 * source_img_mask, dim=-2) / (torch.sum(source_img_mask, dim=-2) + 1e-5)
return mean, var
def forward(self, img_feat_rgb_dir, source_img_mask):
# img_feat_rgb_dir b,d,n,f
b, d, n, _ = img_feat_rgb_dir.shape
view_feat = self.view_fc(img_feat_rgb_dir[..., -4:]) # b,d,n,f-4
img_feat_rgb = img_feat_rgb_dir[..., :-4] + view_feat
mean_feat, var_feat = self.masked_mean_var(img_feat_rgb, source_img_mask)
var_feat = var_feat.view(b, -1, 1, self.feat_ch).repeat(1, 1, n, 1)
avg_feat = mean_feat.view(b, -1, 1, self.feat_ch).repeat(1, 1, n, 1)
feat = torch.cat([img_feat_rgb, var_feat, avg_feat], dim=-1) # b,d,n,f
global_feat = self.global_fc(feat) # b,d,n,f
logits = self.agg_w_fc(global_feat) # b,d,n,1
source_img_mask_ = source_img_mask.reshape(b, 1, n, 1).repeat(1, logits.shape[1], 1, 1) == 0
logits[source_img_mask_] = -1e7
agg_w = F.softmax(logits, dim=-2)
im_feat = (global_feat * agg_w).sum(dim=-2)
return self.fc(im_feat)
================================================
FILE: 3drecon/renderer/cost_reg_net.py
================================================
import torch.nn as nn
class ConvBnReLU3D(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1, norm_act=nn.BatchNorm3d):
super(ConvBnReLU3D, self).__init__()
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
self.bn = norm_act(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
return self.relu(self.bn(self.conv(x)))
class CostRegNet(nn.Module):
def __init__(self, in_channels, norm_act=nn.BatchNorm3d):
super(CostRegNet, self).__init__()
self.conv0 = ConvBnReLU3D(in_channels, 8, norm_act=norm_act)
self.conv1 = ConvBnReLU3D(8, 16, stride=2, norm_act=norm_act)
self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act)
self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act)
self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act)
self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act)
self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act)
self.conv7 = nn.Sequential(
nn.ConvTranspose3d(64, 32, 3, padding=1, output_padding=1, stride=2, bias=False),
norm_act(32)
)
self.conv9 = nn.Sequential(
nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1, stride=2, bias=False),
norm_act(16)
)
self.conv11 = nn.Sequential(
nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1,stride=2, bias=False),
norm_act(8)
)
self.depth_conv = nn.Sequential(nn.Conv3d(8, 1, 3, padding=1, bias=False))
self.feat_conv = nn.Sequential(nn.Conv3d(8, 8, 3, padding=1, bias=False))
def forward(self, x):
conv0 = self.conv0(x)
conv2 = self.conv2(self.conv1(conv0))
conv4 = self.conv4(self.conv3(conv2))
x = self.conv6(self.conv5(conv4))
x = conv4 + self.conv7(x)
del conv4
x = conv2 + self.conv9(x)
del conv2
x = conv0 + self.conv11(x)
del conv0
feat = self.feat_conv(x)
depth = self.depth_conv(x)
return feat, depth
class MinCostRegNet(nn.Module):
def __init__(self, in_channels, norm_act=nn.BatchNorm3d):
super(MinCostRegNet, self).__init__()
self.conv0 = ConvBnReLU3D(in_channels, 8, norm_act=norm_act)
self.conv1 = ConvBnReLU3D(8, 16, stride=2, norm_act=norm_act)
self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act)
self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act)
self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act)
self.conv9 = nn.Sequential(
nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1,
stride=2, bias=False),
norm_act(16))
self.conv11 = nn.Sequential(
nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1,
stride=2, bias=False),
norm_act(8))
self.depth_conv = nn.Sequential(nn.Conv3d(8, 1, 3, padding=1, bias=False))
self.feat_conv = nn.Sequential(nn.Conv3d(8, 8, 3, padding=1, bias=False))
def forward(self, x):
conv0 = self.conv0(x)
conv2 = self.conv2(self.conv1(conv0))
conv4 = self.conv4(self.conv3(conv2))
x = conv4
x = conv2 + self.conv9(x)
del conv2
x = conv0 + self.conv11(x)
del conv0
feat = self.feat_conv(x)
depth = self.depth_conv(x)
return feat, depth
================================================
FILE: 3drecon/renderer/dummy_dataset.py
================================================
import pytorch_lightning as pl
from torch.utils.data import Dataset
import webdataset as wds
from torch.utils.data.distributed import DistributedSampler
class DummyDataset(pl.LightningDataModule):
def __init__(self,seed):
super().__init__()
def setup(self, stage):
if stage in ['fit']:
self.train_dataset = DummyData(True)
self.val_dataset = DummyData(False)
else:
raise NotImplementedError
def train_dataloader(self):
return wds.WebLoader(self.train_dataset, batch_size=1, num_workers=0, shuffle=False)
def val_dataloader(self):
return wds.WebLoader(self.val_dataset, batch_size=1, num_workers=0, shuffle=False)
def test_dataloader(self):
return wds.WebLoader(DummyData(False))
class DummyData(Dataset):
def __init__(self,is_train):
self.is_train=is_train
def __len__(self):
if self.is_train:
return 99999999
else:
return 1
def __getitem__(self, index):
return {}
================================================
FILE: 3drecon/renderer/feature_net.py
================================================
import torch.nn as nn
import torch.nn.functional as F
class ConvBnReLU(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1, norm_act=nn.BatchNorm2d):
super(ConvBnReLU, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
self.bn = norm_act(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
return self.relu(self.bn(self.conv(x)))
class FeatureNet(nn.Module):
def __init__(self, norm_act=nn.BatchNorm2d):
super(FeatureNet, self).__init__()
self.conv0 = nn.Sequential(ConvBnReLU(3, 8, 3, 1, 1, norm_act=norm_act), ConvBnReLU(8, 8, 3, 1, 1, norm_act=norm_act))
self.conv1 = nn.Sequential(ConvBnReLU(8, 16, 5, 2, 2, norm_act=norm_act), ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act))
self.conv2 = nn.Sequential(ConvBnReLU(16, 32, 5, 2, 2, norm_act=norm_act), ConvBnReLU(32, 32, 3, 1, 1, norm_act=norm_act))
self.toplayer = nn.Conv2d(32, 32, 1)
self.lat1 = nn.Conv2d(16, 32, 1)
self.lat0 = nn.Conv2d(8, 32, 1)
self.smooth1 = nn.Conv2d(32, 16, 3, padding=1)
self.smooth0 = nn.Conv2d(32, 8, 3, padding=1)
def _upsample_add(self, x, y):
return F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) + y
def forward(self, x):
conv0 = self.conv0(x)
conv1 = self.conv1(conv0)
conv2 = self.conv2(conv1)
feat2 = self.toplayer(conv2)
feat1 = self._upsample_add(feat2, self.lat1(conv1))
feat0 = self._upsample_add(feat1, self.lat0(conv0))
feat1 = self.smooth1(feat1)
feat0 = self.smooth0(feat0)
return feat2, feat1, feat0
================================================
FILE: 3drecon/renderer/neus_networks.py
================================================
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import tinycudann as tcnn
# Positional encoding embedding. Code was taken from https://github.com/bmild/nerf.
class Embedder:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.create_embedding_fn()
def create_embedding_fn(self):
embed_fns = []
d = self.kwargs['input_dims']
out_dim = 0
if self.kwargs['include_input']:
embed_fns.append(lambda x: x)
out_dim += d
max_freq = self.kwargs['max_freq_log2']
N_freqs = self.kwargs['num_freqs']
if self.kwargs['log_sampling']:
freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
else:
freq_bands = torch.linspace(2. ** 0., 2. ** max_freq, N_freqs)
for freq in freq_bands:
for p_fn in self.kwargs['periodic_fns']:
embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
out_dim += d
self.embed_fns = embed_fns
self.out_dim = out_dim
def embed(self, inputs):
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
def get_embedder(multires, input_dims=3):
embed_kwargs = {
'include_input': True,
'input_dims': input_dims,
'max_freq_log2': multires - 1,
'num_freqs': multires,
'log_sampling': True,
'periodic_fns': [torch.sin, torch.cos],
}
embedder_obj = Embedder(**embed_kwargs)
def embed(x, eo=embedder_obj): return eo.embed(x)
return embed, embedder_obj.out_dim
class SDFNetwork(nn.Module):
def __init__(self, d_in, d_out, d_hidden, n_layers, skip_in=(4,), multires=0, bias=0.5,
scale=1, geometric_init=True, weight_norm=True, inside_outside=False):
super(SDFNetwork, self).__init__()
dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
self.embed_fn_fine = None
if multires > 0:
embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
self.embed_fn_fine = embed_fn
dims[0] = input_ch
self.num_layers = len(dims)
self.skip_in = skip_in
self.scale = scale
for l in range(0, self.num_layers - 1):
if l + 1 in self.skip_in:
out_dim = dims[l + 1] - dims[0]
else:
out_dim = dims[l + 1]
lin = nn.Linear(dims[l], out_dim)
if geometric_init:
if l == self.num_layers - 2:
if not inside_outside:
torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
torch.nn.init.constant_(lin.bias, -bias)
else:
torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
torch.nn.init.constant_(lin.bias, bias)
elif multires > 0 and l == 0:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
elif multires > 0 and l in self.skip_in:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)
else:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
if weight_norm:
lin = nn.utils.weight_norm(lin)
setattr(self, "lin" + str(l), lin)
self.activation = nn.Softplus(beta=100)
def forward(self, inputs):
inputs = inputs * self.scale
if self.embed_fn_fine is not None:
inputs = self.embed_fn_fine(inputs)
x = inputs
for l in range(0, self.num_layers - 1):
lin = getattr(self, "lin" + str(l))
if l in self.skip_in:
x = torch.cat([x, inputs], -1) / np.sqrt(2)
x = lin(x)
if l < self.num_layers - 2:
x = self.activation(x)
return x
def sdf(self, x):
return self.forward(x)[..., :1]
def sdf_hidden_appearance(self, x):
return self.forward(x)
def gradient(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return gradients
def sdf_normal(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return y[..., :1].detach(), gradients.detach()
class SDFNetworkWithFeature(nn.Module):
def __init__(self, cube, dp_in, df_in, d_out, d_hidden, n_layers, skip_in=(4,), multires=0, bias=0.5,
scale=1, geometric_init=True, weight_norm=True, inside_outside=False, cube_length=0.5):
super().__init__()
self.register_buffer("cube", cube)
self.cube_length = cube_length
dims = [dp_in+df_in] + [d_hidden for _ in range(n_layers)] + [d_out]
self.embed_fn_fine = None
if multires > 0:
embed_fn, input_ch = get_embedder(multires, input_dims=dp_in)
self.embed_fn_fine = embed_fn
dims[0] = input_ch + df_in
self.num_layers = len(dims)
self.skip_in = skip_in
self.scale = scale
for l in range(0, self.num_layers - 1):
if l + 1 in self.skip_in:
out_dim = dims[l + 1] - dims[0]
else:
out_dim = dims[l + 1]
lin = nn.Linear(dims[l], out_dim)
if geometric_init:
if l == self.num_layers - 2:
if not inside_outside:
torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
torch.nn.init.constant_(lin.bias, -bias)
else:
torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
torch.nn.init.constant_(lin.bias, bias)
elif multires > 0 and l == 0:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
elif multires > 0 and l in self.skip_in:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)
else:
torch.nn.init.constant_(lin.bias, 0.0)
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
if weight_norm:
lin = nn.utils.weight_norm(lin)
setattr(self, "lin" + str(l), lin)
self.activation = nn.Softplus(beta=100)
def forward(self, points):
points = points * self.scale
# note: point*2 because the cube is [-0.5,0.5]
with torch.no_grad():
feats = F.grid_sample(self.cube, points.view(1,-1,1,1,3)/self.cube_length, mode='bilinear', align_corners=True, padding_mode='zeros').detach()
feats = feats.view(self.cube.shape[1], -1).permute(1,0).view(*points.shape[:-1], -1)
if self.embed_fn_fine is not None:
points = self.embed_fn_fine(points)
x = torch.cat([points, feats], -1)
for l in range(0, self.num_layers - 1):
lin = getattr(self, "lin" + str(l))
if l in self.skip_in:
x = torch.cat([x, points, feats], -1) / np.sqrt(2)
x = lin(x)
if l < self.num_layers - 2:
x = self.activation(x)
# concat feats
x = torch.cat([x, feats], -1)
return x
def sdf(self, x):
return self.forward(x)[..., :1]
def sdf_hidden_appearance(self, x):
return self.forward(x)
def gradient(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return gradients
def sdf_normal(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return y[..., :1].detach(), gradients.detach()
class VanillaMLP(nn.Module):
def __init__(self, dim_in, dim_out, n_neurons, n_hidden_layers):
super().__init__()
self.n_neurons, self.n_hidden_layers = n_neurons, n_hidden_layers
self.sphere_init, self.weight_norm = True, True
self.sphere_init_radius = 0.5
self.layers = [self.make_linear(dim_in, self.n_neurons, is_first=True, is_last=False), self.make_activation()]
for i in range(self.n_hidden_layers - 1):
self.layers += [self.make_linear(self.n_neurons, self.n_neurons, is_first=False, is_last=False), self.make_activation()]
self.layers += [self.make_linear(self.n_neurons, dim_out, is_first=False, is_last=True)]
self.layers = nn.Sequential(*self.layers)
@torch.cuda.amp.autocast(False)
def forward(self, x):
x = self.layers(x.float())
return x
def make_linear(self, dim_in, dim_out, is_first, is_last):
layer = nn.Linear(dim_in, dim_out, bias=True) # network without bias will degrade quality
if self.sphere_init:
if is_last:
torch.nn.init.constant_(layer.bias, -self.sphere_init_radius)
torch.nn.init.normal_(layer.weight, mean=math.sqrt(math.pi) / math.sqrt(dim_in), std=0.0001)
elif is_first:
torch.nn.init.constant_(layer.bias, 0.0)
torch.nn.init.constant_(layer.weight[:, 3:], 0.0)
torch.nn.init.normal_(layer.weight[:, :3], 0.0, math.sqrt(2) / math.sqrt(dim_out))
else:
torch.nn.init.constant_(layer.bias, 0.0)
torch.nn.init.normal_(layer.weight, 0.0, math.sqrt(2) / math.sqrt(dim_out))
else:
torch.nn.init.constant_(layer.bias, 0.0)
torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
if self.weight_norm:
layer = nn.utils.weight_norm(layer)
return layer
def make_activation(self):
if self.sphere_init:
return nn.Softplus(beta=100)
else:
return nn.ReLU(inplace=True)
class SDFHashGridNetwork(nn.Module):
def __init__(self, bound=0.5, feats_dim=13):
super().__init__()
self.bound = bound
# max_resolution = 32
# base_resolution = 16
# n_levels = 4
# log2_hashmap_size = 16
# n_features_per_level = 8
max_resolution = 2048
base_resolution = 16
n_levels = 16
log2_hashmap_size = 19
n_features_per_level = 2
# max_res = base_res * t^(k-1)
per_level_scale = (max_resolution / base_resolution)** (1 / (n_levels - 1))
self.encoder = tcnn.Encoding(
n_input_dims=3,
encoding_config={
"otype": "HashGrid",
"n_levels": n_levels,
"n_features_per_level": n_features_per_level,
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_resolution,
"per_level_scale": per_level_scale,
},
)
self.sdf_mlp = VanillaMLP(n_levels*n_features_per_level+3,feats_dim,64,1)
def forward(self, x):
shape = x.shape[:-1]
x = x.reshape(-1, 3)
x_ = (x + self.bound) / (2 * self.bound)
feats = self.encoder(x_)
feats = torch.cat([x, feats], 1)
feats = self.sdf_mlp(feats)
feats = feats.reshape(*shape,-1)
return feats
def sdf(self, x):
return self(x)[...,:1]
def gradient(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return gradients
def sdf_normal(self, x):
x.requires_grad_(True)
with torch.enable_grad():
y = self.sdf(x)
d_output = torch.ones_like(y, requires_grad=False, device=y.device)
gradients = torch.autograd.grad(
outputs=y,
inputs=x,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
return y[..., :1].detach(), gradients.detach()
class RenderingFFNetwork(nn.Module):
def __init__(self, in_feats_dim=12):
super().__init__()
self.dir_encoder = tcnn.Encoding(
n_input_dims=3,
encoding_config={
"otype": "SphericalHarmonics",
"degree": 4,
},
)
self.color_mlp = tcnn.Network(
n_input_dims = in_feats_dim + 3 + self.dir_encoder.n_output_dims,
n_output_dims = 3,
network_config={
"otype": "FullyFusedMLP",
"activation": "ReLU",
"output_activation": "none",
"n_neurons": 64,
"n_hidden_layers": 2,
},
)
def forward(self, points, normals, view_dirs, feature_vectors):
normals = F.normalize(normals, dim=-1)
view_dirs = F.normalize(view_dirs, dim=-1)
reflective = torch.sum(view_dirs * normals, -1, keepdim=True) * normals * 2 - view_dirs
x = torch.cat([feature_vectors, normals, self.dir_encoder(reflective)], -1)
colors = self.color_mlp(x).float()
colors = F.sigmoid(colors)
return colors
# This implementation is borrowed from IDR: https://github.com/lioryariv/idr
class RenderingNetwork(nn.Module):
def __init__(self, d_feature, d_in, d_out, d_hidden,
n_layers, weight_norm=True, multires_view=0, squeeze_out=True, use_view_dir=True):
super().__init__()
self.squeeze_out = squeeze_out
self.rgb_act=F.sigmoid
self.use_view_dir=use_view_dir
dims = [d_in + d_feature] + [d_hidden for _ in range(n_layers)] + [d_out]
self.embedview_fn = None
if multires_view > 0:
embedview_fn, input_ch = get_embedder(multires_view)
self.embedview_fn = embedview_fn
dims[0] += (input_ch - 3)
self.num_layers = len(dims)
for l in range(0, self.num_layers - 1):
out_dim = dims[l + 1]
lin = nn.Linear(dims[l], out_dim)
if weight_norm:
lin = nn.utils.weight_norm(lin)
setattr(self, "lin" + str(l), lin)
self.relu = nn.ReLU()
def forward(self, points, normals, view_dirs, feature_vectors):
if self.use_view_dir:
view_dirs = F.normalize(view_dirs, dim=-1)
normals = F.normalize(normals, dim=-1)
reflective = torch.sum(view_dirs*normals, -1, keepdim=True) * normals * 2 - view_dirs
if self.embedview_fn is not None: reflective = self.embedview_fn(reflective)
rendering_input = torch.cat([points, reflective, normals, feature_vectors], dim=-1)
else:
rendering_input = torch.cat([points, normals, feature_vectors], dim=-1)
x = rendering_input
for l in range(0, self.num_layers - 1):
lin = getattr(self, "lin" + str(l))
x = lin(x)
if l < self.num_layers - 2:
x = self.relu(x)
if self.squeeze_out:
x = self.rgb_act(x)
return x
class SingleVarianceNetwork(nn.Module):
def __init__(self, init_val, activation='exp'):
super(SingleVarianceNetwork, self).__init__()
self.act = activation
self.register_parameter('variance', nn.Parameter(torch.tensor(init_val)))
def forward(self, x):
device = x.device
if self.act=='exp':
return torch.ones([*x.shape[:-1], 1], dtype=torch.float32, device=device) * torch.exp(self.variance * 10.0)
else:
raise NotImplementedError
def warp(self, x, inv_s):
device = x.device
return torch.ones([*x.shape[:-1], 1], dtype=torch.float32, device=device) * inv_s
================================================
FILE: 3drecon/renderer/ngp_renderer.py
================================================
import math
import trimesh
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from packaging import version as pver
import tinycudann as tcnn
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd
import raymarching
def custom_meshgrid(*args):
# ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid
if pver.parse(torch.__version__) < pver.parse('1.10'):
return torch.meshgrid(*args)
else:
return torch.meshgrid(*args, indexing='ij')
def sample_pdf(bins, weights, n_samples, det=False):
# This implementation is from NeRF
# bins: [B, T], old_z_vals
# weights: [B, T - 1], bin weights.
# return: [B, n_samples], new_z_vals
# Get pdf
weights = weights + 1e-5 # prevent nans
pdf = weights / torch.sum(weights, -1, keepdim=True)
cdf = torch.cumsum(pdf, -1)
cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
# Take uniform samples
if det:
u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples).to(weights.device)
u = u.expand(list(cdf.shape[:-1]) + [n_samples])
else:
u = torch.rand(list(cdf.shape[:-1]) + [n_samples]).to(weights.device)
# Invert CDF
u = u.contiguous()
inds = torch.searchsorted(cdf, u, right=True)
below = torch.max(torch.zeros_like(inds - 1), inds - 1)
above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
inds_g = torch.stack([below, above], -1) # (B, n_samples, 2)
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
denom = (cdf_g[..., 1] - cdf_g[..., 0])
denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
return samples
def plot_pointcloud(pc, color=None):
# pc: [N, 3]
# color: [N, 3/4]
print('[visualize points]', pc.shape, pc.dtype, pc.min(0), pc.max(0))
pc = trimesh.PointCloud(pc, color)
# axis
axes = trimesh.creation.axis(axis_length=4)
# sphere
sphere = trimesh.creation.icosphere(radius=1)
trimesh.Scene([pc, axes, sphere]).show()
class NGPRenderer(nn.Module):
def __init__(self,
bound=1,
cuda_ray=True,
density_scale=1, # scale up deltas (or sigmas), to make the density grid more sharp. larger value than 1 usually improves performance.
min_near=0.2,
density_thresh=0.01,
bg_radius=-1,
):
super().__init__()
self.bound = bound
self.cascade = 1
self.grid_size = 128
self.density_scale = density_scale
self.min_near = min_near
self.density_thresh = density_thresh
self.bg_radius = bg_radius # radius of the background sphere.
# prepare aabb with a 6D tensor (xmin, ymin, zmin, xmax, ymax, zmax)
# NOTE: aabb (can be rectangular) is only used to generate points, we still rely on bound (always cubic) to calculate density grid and hashing.
aabb_train = torch.FloatTensor([-bound, -bound, -bound, bound, bound, bound])
aabb_infer = aabb_train.clone()
self.register_buffer('aabb_train', aabb_train)
self.register_buffer('aabb_infer', aabb_infer)
# extra state for cuda raymarching
self.cuda_ray = cuda_ray
if cuda_ray:
# density grid
density_grid = torch.zeros([self.cascade, self.grid_size ** 3]) # [CAS, H * H * H]
density_bitfield = torch.zeros(self.cascade * self.grid_size ** 3 // 8, dtype=torch.uint8) # [CAS * H * H * H // 8]
self.register_buffer('density_grid', density_grid)
self.register_buffer('density_bitfield', density_bitfield)
self.mean_density = 0
self.iter_density = 0
# step counter
step_counter = torch.zeros(16, 2, dtype=torch.int32) # 16 is hardcoded for averaging...
self.register_buffer('step_counter', step_counter)
self.mean_count = 0
self.local_step = 0
def forward(self, x, d):
raise NotImplementedError()
# separated density and color query (can accelerate non-cuda-ray mode.)
def density(self, x):
raise NotImplementedError()
def color(self, x, d, mask=None, **kwargs):
raise NotImplementedError()
def reset_extra_state(self):
if not self.cuda_ray:
return
# density grid
self.density_grid.zero_()
self.mean_density = 0
self.iter_density = 0
# step counter
self.step_counter.zero_()
self.mean_count = 0
self.local_step = 0
def run(self, rays_o, rays_d, num_steps=128, upsample_steps=128, bg_color=None, perturb=False, **kwargs):
# rays_o, rays_d: [B, N, 3], assumes B == 1
# bg_color: [3] in range [0, 1]
# return: image: [B, N, 3], depth: [B, N]
prefix = rays_o.shape[:-1]
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
N = rays_o.shape[0] # N = B * N, in fact
device = rays_o.device
# choose aabb
aabb = self.aabb_train if self.training else self.aabb_infer
# sample steps
nears, fars = raymarching.near_far_from_aabb(rays_o, rays_d, aabb, self.min_near)
nears.unsqueeze_(-1)
fars.unsqueeze_(-1)
#print(f'nears = {nears.min().item()} ~ {nears.max().item()}, fars = {fars.min().item()} ~ {fars.max().item()}')
z_vals = torch.linspace(0.0, 1.0, num_steps, device=device).unsqueeze(0) # [1, T]
z_vals = z_vals.expand((N, num_steps)) # [N, T]
z_vals = nears + (fars - nears) * z_vals # [N, T], in [nears, fars]
# perturb z_vals
sample_dist = (fars - nears) / num_steps
if perturb:
z_vals = z_vals + (torch.rand(z_vals.shape, device=device) - 0.5) * sample_dist
#z_vals = z_vals.clamp(nears, fars) # avoid out of bounds xyzs.
# generate xyzs
xyzs = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * z_vals.unsqueeze(-1) # [N, 1, 3] * [N, T, 1] -> [N, T, 3]
xyzs = torch.min(torch.max(xyzs, aabb[:3]), aabb[3:]) # a manual clip.
#plot_pointcloud(xyzs.reshape(-1, 3).detach().cpu().numpy())
# query SDF and RGB
density_outputs = self.density(xyzs.reshape(-1, 3))
#sigmas = density_outputs['sigma'].view(N, num_steps) # [N, T]
for k, v in density_outputs.items():
density_outputs[k] = v.view(N, num_steps, -1)
# upsample z_vals (nerf-like)
if upsample_steps > 0:
with torch.no_grad():
deltas = z_vals[..., 1:] - z_vals[..., :-1] # [N, T-1]
deltas = torch.cat([deltas, sample_dist * torch.ones_like(deltas[..., :1])], dim=-1)
alphas = 1 - torch.exp(-deltas * self.density_scale * density_outputs['sigma'].squeeze(-1)) # [N, T]
alphas_shifted = torch.cat([torch.ones_like(alphas[..., :1]), 1 - alphas + 1e-15], dim=-1) # [N, T+1]
weights = alphas * torch.cumprod(alphas_shifted, dim=-1)[..., :-1] # [N, T]
# sample new z_vals
z_vals_mid = (z_vals[..., :-1] + 0.5 * deltas[..., :-1]) # [N, T-1]
new_z_vals = sample_pdf(z_vals_mid, weights[:, 1:-1], upsample_steps, det=not self.training).detach() # [N, t]
new_xyzs = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * new_z_vals.unsqueeze(-1) # [N, 1, 3] * [N, t, 1] -> [N, t, 3]
new_xyzs = torch.min(torch.max(new_xyzs, aabb[:3]), aabb[3:]) # a manual clip.
# only forward new points to save computation
new_density_outputs = self.density(new_xyzs.reshape(-1, 3))
#new_sigmas = new_density_outputs['sigma'].view(N, upsample_steps) # [N, t]
for k, v in new_density_outputs.items():
new_density_outputs[k] = v.view(N, upsample_steps, -1)
# re-order
z_vals = torch.cat([z_vals, new_z_vals], dim=1) # [N, T+t]
z_vals, z_index = torch.sort(z_vals, dim=1)
xyzs = torch.cat([xyzs, new_xyzs], dim=1) # [N, T+t, 3]
xyzs = torch.gather(xyzs, dim=1, index=z_index.unsqueeze(-1).expand_as(xyzs))
for k in density_outputs:
tmp_output = torch.cat([density_outputs[k], new_density_outputs[k]], dim=1)
density_outputs[k] = torch.gather(tmp_output, dim=1, index=z_index.unsqueeze(-1).expand_as(tmp_output))
deltas = z_vals[..., 1:] - z_vals[..., :-1] # [N, T+t-1]
deltas = torch.cat([deltas, sample_dist * torch.ones_like(deltas[..., :1])], dim=-1)
alphas = 1 - torch.exp(-deltas * self.density_scale * density_outputs['sigma'].squeeze(-1)) # [N, T+t]
alphas_shifted = torch.cat([torch.ones_like(alphas[..., :1]), 1 - alphas + 1e-15], dim=-1) # [N, T+t+1]
weights = alphas * torch.cumprod(alphas_shifted, dim=-1)[..., :-1] # [N, T+t]
dirs = rays_d.view(-1, 1, 3).expand_as(xyzs)
for k, v in density_outputs.items():
density_outputs[k] = v.view(-1, v.shape[-1])
mask = weights > 1e-4 # hard coded
rgbs = self.color(xyzs.reshape(-1, 3), dirs.reshape(-1, 3), mask=mask.reshape(-1), **density_outputs)
rgbs = rgbs.view(N, -1, 3) # [N, T+t, 3]
#print(xyzs.shape, 'valid_rgb:', mask.sum().item())
# calculate weight_sum (mask)
weights_sum = weights.sum(dim=-1) # [N]
# calculate depth
ori_z_vals = ((z_vals - nears) / (fars - nears)).clamp(0, 1)
depth = torch.sum(weights * ori_z_vals, dim=-1)
# calculate color
image = torch.sum(weights.unsqueeze(-1) * rgbs, dim=-2) # [N, 3], in [0, 1]
# mix background color
if self.bg_radius > 0:
# use the bg model to calculate bg_color
sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
bg_color = self.background(sph, rays_d.reshape(-1, 3)) # [N, 3]
elif bg_color is None:
bg_color = 1
image = image + (1 - weights_sum).unsqueeze(-1) * bg_color
image = image.view(*prefix, 3)
depth = depth.view(*prefix)
# tmp: reg loss in mip-nerf 360
# z_vals_shifted = torch.cat([z_vals[..., 1:], sample_dist * torch.ones_like(z_vals[..., :1])], dim=-1)
# mid_zs = (z_vals + z_vals_shifted) / 2 # [N, T]
# loss_dist = (torch.abs(mid_zs.unsqueeze(1) - mid_zs.unsqueeze(2)) * (weights.unsqueeze(1) * weights.unsqueeze(2))).sum() + 1/3 * ((z_vals_shifted - z_vals_shifted) * (weights ** 2)).sum()
return {
'depth': depth,
'image': image,
'weights_sum': weights_sum,
}
def run_cuda(self, rays_o, rays_d, dt_gamma=0, bg_color=None, perturb=False, force_all_rays=False, max_steps=1024, T_thresh=1e-4, **kwargs):
# rays_o, rays_d: [B, N, 3], assumes B == 1
# return: image: [B, N, 3], depth: [B, N]
prefix = rays_o.shape[:-1]
rays_o = rays_o.contiguous().view(-1, 3)
rays_d = rays_d.contiguous().view(-1, 3)
N = rays_o.shape[0] # N = B * N, in fact
device = rays_o.device
# pre-calculate near far
nears, fars = raymarching.near_far_from_aabb(rays_o, rays_d, self.aabb_train if self.training else self.aabb_infer, self.min_near)
# mix background color
if self.bg_radius > 0:
# use the bg model to calculate bg_color
sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
bg_color = self.background(sph, rays_d) # [N, 3]
elif bg_color is None:
bg_color = 1
results = {}
if self.training:
# setup counter
counter = self.step_counter[self.local_step % 16]
counter.zero_() # set to 0
self.local_step += 1
xyzs, dirs, deltas, rays = raymarching.march_rays_train(rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, counter, self.mean_count, perturb, 128, force_all_rays, dt_gamma, max_steps)
#plot_pointcloud(xyzs.reshape(-1, 3).detach().cpu().numpy())
sigmas, rgbs = self(xyzs, dirs)
sigmas = self.density_scale * sigmas
weights_sum, depth, image = raymarching.composite_rays_train(sigmas, rgbs, deltas, rays, T_thresh)
image = image + (1 - weights_sum).unsqueeze(-1) * bg_color
depth = torch.clamp(depth - nears, min=0) / (fars - nears)
image = image.view(*prefix, 3)
depth = depth.view(*prefix)
else:
# allocate outputs
# if use autocast, must init as half so it won't be autocasted and lose reference.
#dtype = torch.half if torch.is_autocast_enabled() else torch.float32
# output should always be float32! only network inference uses half.
dtype = torch.float32
weights_sum = torch.zeros(N, dtype=dtype, device=device)
depth = torch.zeros(N, dtype=dtype, device=device)
image = torch.zeros(N, 3, dtype=dtype, device=device)
n_alive = N
rays_alive = torch.arange(n_alive, dtype=torch.int32, device=device) # [N]
rays_t = nears.clone() # [N]
step = 0
while step < max_steps:
# count alive rays
n_alive = rays_alive.shape[0]
# exit loop
if n_alive <= 0:
break
# decide compact_steps
n_step = max(min(N // n_alive, 8), 1)
xyzs, dirs, deltas = raymarching.march_rays(n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, self.bound, self.density_bitfield, self.cascade, self.grid_size, nears, fars, 128, perturb if step == 0 else False, dt_gamma, max_steps)
sigmas, rgbs = self(xyzs, dirs)
# density_outputs = self.density(xyzs) # [M,], use a dict since it may include extra things, like geo_feat for rgb.
# sigmas = density_outputs['sigma']
# rgbs = self.color(xyzs, dirs, **density_outputs)
sigmas = self.density_scale * sigmas
raymarching.composite_rays(n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, deltas, weights_sum, depth, image, T_thresh)
rays_alive = rays_alive[rays_alive >= 0]
#print(f'step = {step}, n_step = {n_step}, n_alive = {n_alive}, xyzs: {xyzs.shape}')
step += n_step
image = image + (1 - weights_sum).unsqueeze(-1) * bg_color
depth = torch.clamp(depth - nears, min=0) / (fars - nears)
image = image.view(*prefix, 3)
depth = depth.view(*prefix)
results['weights_sum'] = weights_sum
results['depth'] = depth
results['image'] = image
return results
@torch.no_grad()
def mark_untrained_grid(self, poses, intrinsic, S=64):
# poses: [B, 4, 4]
# intrinsic: [3, 3]
if not self.cuda_ray:
return
if isinstance(poses, np.ndarray):
poses = torch.from_numpy(poses)
B = poses.shape[0]
fx, fy, cx, cy = intrinsic
X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
count = torch.zeros_like(self.density_grid)
poses = poses.to(count.device)
# 5-level loop, forgive me...
for xs in X:
for ys in Y:
for zs in Z:
# construct points
xx, yy, zz = custom_meshgrid(xs, ys, zs)
coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128)
indices = raymarching.morton3D(coords).long() # [N]
world_xyzs = (2 * coords.float() / (self.grid_size - 1) - 1).unsqueeze(0) # [1, N, 3] in [-1, 1]
# cascading
for cas in range(self.cascade):
bound = min(2 ** cas, self.bound)
half_grid_size = bound / self.grid_size
# scale to current cascade's resolution
cas_world_xyzs = world_xyzs * (bound - half_grid_size)
# split batch to avoid OOM
head = 0
while head < B:
tail = min(head + S, B)
# world2cam transform (poses is c2w, so we need to transpose it. Another transpose is needed for batched matmul, so the final form is without transpose.)
cam_xyzs = cas_world_xyzs - poses[head:tail, :3, 3].unsqueeze(1)
cam_xyzs = cam_xyzs @ poses[head:tail, :3, :3] # [S, N, 3]
# query if point is covered by any camera
mask_z = cam_xyzs[:, :, 2] > 0 # [S, N]
mask_x = torch.abs(cam_xyzs[:, :, 0]) < cx / fx * cam_xyzs[:, :, 2] + half_grid_size * 2
mask_y = torch.abs(cam_xyzs[:, :, 1]) < cy / fy * cam_xyzs[:, :, 2] + half_grid_size * 2
mask = (mask_z & mask_x & mask_y).sum(0).reshape(-1) # [N]
# update count
count[cas, indices] += mask
head += S
# mark untrained grid as -1
self.density_grid[count == 0] = -1
print(f'[mark untrained grid] {(count == 0).sum()} from {self.grid_size ** 3 * self.cascade}')
@torch.no_grad()
def update_extra_state(self, decay=0.95, S=128):
# call before each epoch to update extra states.
if not self.cuda_ray:
return
### update density grid
tmp_grid = - torch.ones_like(self.density_grid)
# full update.
if self.iter_density < 16:
#if True:
X = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
Y = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
Z = torch.arange(self.grid_size, dtype=torch.int32, device=self.density_bitfield.device).split(S)
for xs in X:
for ys in Y:
for zs in Z:
# construct points
xx, yy, zz = custom_meshgrid(xs, ys, zs)
coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [N, 3], in [0, 128)
indices = raymarching.morton3D(coords).long() # [N]
xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1]
# cascading
for cas in range(self.cascade):
bound = min(2 ** cas, self.bound)
half_grid_size = bound / self.grid_size
# scale to current cascade's resolution
cas_xyzs = xyzs * (bound - half_grid_size)
# add noise in [-hgs, hgs]
cas_xyzs += (torch.rand_like(cas_xyzs) * 2 - 1) * half_grid_size
# query density
sigmas = self.density(cas_xyzs)['sigma'].reshape(-1).detach()
sigmas *= self.density_scale
# assign
tmp_grid[cas, indices] = sigmas
# partial update (half the computation)
# TODO: why no need of maxpool ?
else:
N = self.grid_size ** 3 // 4 # H * H * H / 4
for cas in range(self.cascade):
# random sample some positions
coords = torch.randint(0, self.grid_size, (N, 3), device=self.density_bitfield.device) # [N, 3], in [0, 128)
indices = raymarching.morton3D(coords).long() # [N]
# random sample occupied positions
occ_indices = torch.nonzero(self.density_grid[cas] > 0).squeeze(-1) # [Nz]
rand_mask = torch.randint(0, occ_indices.shape[0], [N], dtype=torch.long, device=self.density_bitfield.device)
occ_indices = occ_indices[rand_mask] # [Nz] --> [N], allow for duplication
occ_coords = raymarching.morton3D_invert(occ_indices) # [N, 3]
# concat
indices = torch.cat([indices, occ_indices], dim=0)
coords = torch.cat([coords, occ_coords], dim=0)
# same below
xyzs = 2 * coords.float() / (self.grid_size - 1) - 1 # [N, 3] in [-1, 1]
bound = min(2 ** cas, self.bound)
half_grid_size = bound / self.grid_size
# scale to current cascade's resolution
cas_xyzs = xyzs * (bound - half_grid_size)
# add noise in [-hgs, hgs]
cas_xyzs += (torch.rand_like(cas_xyzs) * 2 - 1) * half_grid_size
# query density
sigmas = self.density(cas_xyzs)['sigma'].reshape(-1).detach()
sigmas *= self.density_scale
# assign
tmp_grid[cas, indices] = sigmas
## max-pool on tmp_grid for less aggressive culling [No significant improvement...]
# invalid_mask = tmp_grid < 0
# tmp_grid = F.max_pool3d(tmp_grid.view(self.cascade, 1, self.grid_size, self.grid_size, self.grid_size), kernel_size=3, stride=1, padding=1).view(self.cascade, -1)
# tmp_grid[invalid_mask] = -1
# ema update
valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0)
self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask])
self.mean_density = torch.mean(self.density_grid.clamp(min=0)).item() # -1 regions are viewed as 0 density.
#self.mean_density = torch.mean(self.density_grid[self.density_grid > 0]).item() # do not count -1 regions
self.iter_density += 1
# convert to bitfield
density_thresh = min(self.mean_density, self.density_thresh)
self.density_bitfield = raymarching.packbits(self.density_grid, density_thresh, self.density_bitfield)
### update step counter
total_step = min(16, self.local_step)
if total_step > 0:
self.mean_count = int(self.step_counter[:total_step, 0].sum().item() / total_step)
self.local_step = 0
#print(f'[density grid] min={self.density_grid.min().item():.4f}, max={self.density_grid.max().item():.4f}, mean={self.mean_density:.4f}, occ_rate={(self.density_grid > 0.01).sum() / (128**3 * self.cascade):.3f} | [step counter] mean={self.mean_count}')
def render(self, rays_o, rays_d, staged=False, max_ray_batch=4096, **kwargs):
# rays_o, rays_d: [B, N, 3], assumes B == 1
# return: pred_rgb: [B, N, 3]
if self.cuda_ray:
_run = self.run_cuda
else:
_run = self.run
results = _run(rays_o, rays_d, **kwargs)
return results
class _trunc_exp(Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32) # cast to float32
def forward(ctx, x):
ctx.save_for_backward(x)
return torch.exp(x)
@staticmethod
@custom_bwd
def backward(ctx, g):
x = ctx.saved_tensors[0]
return g * torch.exp(x.clamp(-15, 15))
trunc_exp = _trunc_exp.apply
class NGPNetwork(NGPRenderer):
def __init__(self,
num_layers=2,
hidden_dim=64,
geo_feat_dim=15,
num_layers_color=3,
hidden_dim_color=64,
bound=0.5,
max_resolution=128,
base_resolution=16,
n_levels=16,
**kwargs
):
super().__init__(bound, **kwargs)
# sigma network
self.num_layers = num_layers
self.hidden_dim = hidden_dim
self.geo_feat_dim = geo_feat_dim
self.bound = bound
log2_hashmap_size = 19
n_features_per_level = 2
per_level_scale = np.exp2(np.log2(max_resolution / base_resolution) / (n_levels - 1))
self.encoder = tcnn.Encoding(
n_input_dims=3,
encoding_config={
"otype": "HashGrid",
"n_levels": n_levels,
"n_features_per_level": n_features_per_level,
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_resolution,
"per_level_scale": per_level_scale,
},
)
self.sigma_net = tcnn.Network(
n_input_dims = n_levels * 2,
n_output_dims=1 + self.geo_feat_dim,
network_config={
"otype": "FullyFusedMLP",
"activation": "ReLU",
"output_activation": "None",
"n_neurons": hidden_dim,
"n_hidden_layers": num_layers - 1,
},
)
# color network
self.num_layers_color = num_layers_color
self.hidden_dim_color = hidden_dim_color
self.encoder_dir = tcnn.Encoding(
n_input_dims=3,
encoding_config={
"otype": "SphericalHarmonics",
"degree": 4,
},
)
self.in_dim_color = self.encoder_dir.n_output_dims + self.geo_feat_dim
self.color_net = tcnn.Network(
n_input_dims = self.in_dim_color,
n_output_dims=3,
network_config={
"otype": "FullyFusedMLP",
"activation": "ReLU",
"output_activation": "None",
"n_neurons": hidden_dim_color,
"n_hidden_layers": num_layers_color - 1,
},
)
self.density_scale, self.density_std = 10.0, 0.25
def forward(self, x, d):
# x: [N, 3], in [-bound, bound]
# d: [N, 3], nomalized in [-1, 1]
# sigma
x_raw = x
x = (x + self.bound) / (2 * self.bound) # to [0, 1]
x = self.encoder(x)
h = self.sigma_net(x)
# sigma = F.relu(h[..., 0])
density = h[..., 0]
# add density bias
dist = torch.norm(x_raw, dim=-1)
density_bias = (1 - dist / self.density_std) * self.density_scale
density = density_bias + density
sigma = F.softplus(density)
geo_feat = h[..., 1:]
# color
d = (d + 1) / 2 # tcnn SH encoding requires inputs to be in [0, 1]
d = self.encoder_dir(d)
# p = torch.zeros_like(geo_feat[..., :1]) # manual input padding
h = torch.cat([d, geo_feat], dim=-1)
h = self.color_net(h)
# sigmoid activation for rgb
color = torch.sigmoid(h)
return sigma, color
def density(self, x):
# x: [N, 3], in [-bound, bound]
x_raw = x
x = (x + self.bound) / (2 * self.bound) # to [0, 1]
x = self.encoder(x)
h = self.sigma_net(x)
# sigma = F.relu(h[..., 0])
density = h[..., 0]
# add density bias
dist = torch.norm(x_raw, dim=-1)
density_bias = (1 - dist / self.density_std) * self.density_scale
density = density_bias + density
sigma = F.softplus(density)
geo_feat = h[..., 1:]
return {
'sigma': sigma,
'geo_feat': geo_feat,
}
# allow masked inference
def color(self, x, d, mask=None, geo_feat=None, **kwargs):
# x: [N, 3] in [-bound, bound]
# mask: [N,], bool, indicates where we actually needs to compute rgb.
x = (x + self.bound) / (2 * self.bound) # to [0, 1]
if mask is not None:
rgbs = torch.zeros(mask.shape[0], 3, dtype=x.dtype, device=x.device) # [N, 3]
# in case of empty mask
if not mask.any():
return rgbs
x = x[mask]
d = d[mask]
geo_feat = geo_feat[mask]
# color
d = (d + 1) / 2 # tcnn SH encoding requires inputs to be in [0, 1]
d = self.encoder_dir(d)
h = torch.cat([d, geo_feat], dim=-1)
h = self.color_net(h)
# sigmoid activation for rgb
h = torch.sigmoid(h)
if mask is not None:
rgbs[mask] = h.to(rgbs.dtype) # fp16 --> fp32
else:
rgbs = h
return rgbs
================================================
FILE: 3drecon/renderer/renderer.py
================================================
import abc
import os
from pathlib import Path
import cv2
import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from omegaconf import OmegaConf
from skimage.io import imread, imsave
from PIL import Image
from torch.optim.lr_scheduler import LambdaLR
from renderer.neus_networks import SDFNetwork, RenderingNetwork, SingleVarianceNetwork, SDFHashGridNetwork, RenderingFFNetwork
from renderer.ngp_renderer import NGPNetwork
from util import instantiate_from_config, read_pickle, concat_images_list
DEFAULT_RADIUS = np.sqrt(3)/2
DEFAULT_SIDE_LENGTH = 0.6
def sample_pdf(bins, weights, n_samples, det=True):
device = bins.device
dtype = bins.dtype
# This implementation is from NeRF
# Get pdf
weights = weights + 1e-5 # prevent nans
pdf = weights / torch.sum(weights, -1, keepdim=True)
cdf = torch.cumsum(pdf, -1)
cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
# Take uniform samples
if det:
u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples, dtype=dtype, device=device)
u = u.expand(list(cdf.shape[:-1]) + [n_samples])
else:
u = torch.rand(list(cdf.shape[:-1]) + [n_samples], dtype=dtype, device=device)
# Invert CDF
u = u.contiguous()
inds = torch.searchsorted(cdf, u, right=True)
below = torch.max(torch.zeros_like(inds - 1), inds - 1)
above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
denom = (cdf_g[..., 1] - cdf_g[..., 0])
denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
return samples
def near_far_from_sphere(rays_o, rays_d, radius=DEFAULT_RADIUS):
a = torch.sum(rays_d ** 2, dim=-1, keepdim=True)
b = torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
mid = -b / a
near = mid - radius
far = mid + radius
return near, far
class BackgroundRemoval:
def __init__(self, device='cuda'):
from carvekit.api.high import HiInterface
self.interface = HiInterface(
object_type="object", # Can be "object" or "hairs-like".
batch_size_seg=5,
batch_size_matting=1,
device=device,
seg_mask_size=640, # Use 640 for Tracer B7 and 320 for U2Net
matting_mask_size=2048,
trimap_prob_threshold=231,
trimap_dilation=30,
trimap_erosion_iters=5,
fp16=True,
)
@torch.no_grad()
def __call__(self, image):
# image: [H, W, 3] array in [0, 255].
image = Image.fromarray(image)
image = self.interface([image])[0]
image = np.array(image)
return image
class BaseRenderer(nn.Module):
def __init__(self, train_batch_num, test_batch_num):
super().__init__()
self.train_batch_num = train_batch_num
self.test_batch_num = test_batch_num
@abc.abstractmethod
def render_impl(self, ray_batch, is_train, step):
pass
@abc.abstractmethod
def render_with_loss(self, ray_batch, is_train, step):
pass
def render(self, ray_batch, is_train, step):
batch_num = self.train_batch_num if is_train else self.test_batch_num
ray_num = ray_batch['rays_o'].shape[0]
outputs = {}
for ri in range(0, ray_num, batch_num):
cur_ray_batch = {}
for k, v in ray_batch.items():
cur_ray_batch[k] = v[ri:ri + batch_num]
cur_outputs = self.render_impl(cur_ray_batch, is_train, step)
for k, v in cur_outputs.items():
if k not in outputs: outputs[k] = []
outputs[k].append(v)
for k, v in outputs.items():
outputs[k] = torch.cat(v, 0)
return outputs
class NeuSRenderer(BaseRenderer):
def __init__(self, train_batch_num, test_batch_num, lambda_eikonal_loss=0.1, use_mask=True,
lambda_rgb_loss=1.0, lambda_mask_loss=0.0, rgb_loss='soft_l1', coarse_sn=64, fine_sn=64):
super().__init__(train_batch_num, test_batch_num)
self.n_samples = coarse_sn
self.n_importance = fine_sn
self.up_sample_steps = 4
self.anneal_end = 200
self.use_mask = use_mask
self.lambda_eikonal_loss = lambda_eikonal_loss
self.lambda_rgb_loss = lambda_rgb_loss
self.lambda_mask_loss = lambda_mask_loss
self.rgb_loss = rgb_loss
self.sdf_network = SDFNetwork(d_out=257, d_in=3, d_hidden=256, n_layers=8, skip_in=[4], multires=6, bias=0.5, scale=1.0, geometric_init=True, weight_norm=True)
self.color_network = RenderingNetwork(d_feature=256, d_in=9, d_out=3, d_hidden=256, n_layers=4, weight_norm=True, multires_view=4, squeeze_out=True)
self.default_dtype = torch.float32
self.deviation_network = SingleVarianceNetwork(0.3)
@torch.no_grad()
def get_vertex_colors(self, vertices):
"""
@param vertices: n,3
@return:
"""
V = vertices.shape[0]
bn = 20480
verts_colors = []
with torch.no_grad():
for vi in range(0, V, bn):
verts = torch.from_numpy(vertices[vi:vi+bn].astype(np.float32)).cuda()
feats = self.sdf_network(verts)[..., 1:]
gradients = self.sdf_network.gradient(verts) # ...,3
gradients = F.normalize(gradients, dim=-1)
colors = self.color_network(verts, gradients, gradients, feats)
colors = torch.clamp(colors,min=0,max=1).cpu().numpy()
verts_colors.append(colors)
verts_colors = (np.concatenate(verts_colors, 0)*255).astype(np.uint8)
return verts_colors
def upsample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
"""
Up sampling give a fixed inv_s
"""
device = rays_o.device
batch_size, n_samples = z_vals.shape
pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None] # n_rays, n_samples, 3
inner_mask = self.get_inner_mask(pts)
# radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
inside_sphere = inner_mask[:, :-1] | inner_mask[:, 1:]
sdf = sdf.reshape(batch_size, n_samples)
prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
mid_sdf = (prev_sdf + next_sdf) * 0.5
cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
prev_cos_val = torch.cat([torch.zeros([batch_size, 1], dtype=self.default_dtype, device=device), cos_val[:, :-1]], dim=-1)
cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
dist = (next_z_vals - prev_z_vals)
prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
next_esti_sdf = mid_sdf + cos_val * dist * 0.5
prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
weights = alpha * torch.cumprod(
torch.cat([torch.ones([batch_size, 1], dtype=self.default_dtype, device=device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
return z_samples
def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False):
batch_size, n_samples = z_vals.shape
_, n_importance = new_z_vals.shape
pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
z_vals, index = torch.sort(z_vals, dim=-1)
if not last:
device = pts.device
new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
sdf = torch.cat([sdf, new_sdf], dim=-1)
xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1).to(device)
index = index.reshape(-1)
sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
return z_vals, sdf
def sample_depth(self, rays_o, rays_d, near, far, perturb):
n_samples = self.n_samples
n_importance = self.n_importance
up_sample_steps = self.up_sample_steps
device = rays_o.device
# sample points
batch_size = len(rays_o)
z_vals = torch.linspace(0.0, 1.0, n_samples, dtype=self.default_dtype, device=device) # sn
z_vals = near + (far - near) * z_vals[None, :] # rn,sn
if perturb > 0:
t_rand = (torch.rand([batch_size, 1]).to(device) - 0.5)
z_vals = z_vals + t_rand * 2.0 / n_samples
# Up sample
with torch.no_grad():
pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
sdf = self.sdf_network.sdf(pts).reshape(batch_size, n_samples)
for i in range(up_sample_steps):
rn, sn = z_vals.shape
inv_s = torch.ones(rn, sn - 1, dtype=self.default_dtype, device=device) * 64 * 2 ** i
new_z_vals = self.upsample(rays_o, rays_d, z_vals, sdf, n_importance // up_sample_steps, inv_s)
z_vals, sdf = self.cat_z_vals(rays_o, rays_d, z_vals, new_z_vals, sdf, last=(i + 1 == up_sample_steps))
return z_vals
def compute_sdf_alpha(self, points, dists, dirs, cos_anneal_ratio, step):
# points [...,3] dists [...] dirs[...,3]
sdf_nn_output = self.sdf_network(points)
sdf = sdf_nn_output[..., 0]
feature_vector = sdf_nn_output[..., 1:]
gradients = self.sdf_network.gradient(points) # ...,3
inv_s = self.deviation_network(points).clip(1e-6, 1e6) # ...,1
inv_s = inv_s[..., 0]
true_cos = (dirs * gradients).sum(-1) # [...]
iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
F.relu(-true_cos) * cos_anneal_ratio) # always non-positive
# Estimate signed distances at section points
estimated_next_sdf = sdf + iter_cos * dists * 0.5
estimated_prev_sdf = sdf - iter_cos * dists * 0.5
prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
p = prev_cdf - next_cdf
c = prev_cdf
alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0) # [...]
return alpha, gradients, feature_vector, inv_s, sdf
def get_anneal_val(self, step):
if self.anneal_end < 0:
return 1.0
else:
return np.min([1.0, step / self.anneal_end])
def get_inner_mask(self, points):
return torch.sum(torch.abs(points)<=DEFAULT_SIDE_LENGTH,-1)==3
def render_impl(self, ray_batch, is_train, step):
near, far = near_far_from_sphere(ray_batch['rays_o'], ray_batch['rays_d'])
rays_o, rays_d = ray_batch['rays_o'], ray_batch['rays_d']
z_vals = self.sample_depth(rays_o, rays_d, near, far, is_train)
batch_size, n_samples = z_vals.shape
# section length in original space
dists = z_vals[..., 1:] - z_vals[..., :-1] # rn,sn-1
dists = torch.cat([dists, dists[..., -1:]], -1) # rn,sn
mid_z_vals = z_vals + dists * 0.5
points = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * mid_z_vals.unsqueeze(-1) # rn, sn, 3
inner_mask = self.get_inner_mask(points)
dirs = rays_d.unsqueeze(-2).expand(batch_size, n_samples, 3)
dirs = F.normalize(dirs, dim=-1)
device = rays_o.device
alpha, sampled_color, gradient_error, normal = torch.zeros(batch_size, n_samples, dtype=self.default_dtype, device=device), \
torch.zeros(batch_size, n_samples, 3, dtype=self.default_dtype, device=device), \
torch.zeros([batch_size, n_samples], dtype=self.default_dtype, device=device), \
torch.zeros([batch_size, n_samples, 3], dtype=self.default_dtype, device=device)
if torch.sum(inner_mask) > 0:
cos_anneal_ratio = self.get_anneal_val(step) if is_train else 1.0
alpha[inner_mask], gradients, feature_vector, inv_s, sdf = self.compute_sdf_alpha(points[inner_mask], dists[inner_mask], dirs[inner_mask], cos_anneal_ratio, step)
sampled_color[inner_mask] = self.color_network(points[inner_mask], gradients, -dirs[inner_mask], feature_vector)
# Eikonal loss
gradient_error[inner_mask] = (torch.linalg.norm(gradients, ord=2, dim=-1) - 1.0) ** 2 # rn,sn
normal[inner_mask] = F.normalize(gradients, dim=-1)
weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1], dtype=self.default_dtype, device=device), 1. - alpha + 1e-7], -1), -1)[..., :-1] # rn,sn
mask = torch.sum(weights,dim=1).unsqueeze(-1) # rn,1
color = (sampled_color * weights[..., None]).sum(dim=1) + (1 - mask) # add white background
normal = (normal * weights[..., None]).sum(dim=1)
outputs = {
'rgb': color, # rn,3
'gradient_error': gradient_error, # rn,sn
'inner_mask': inner_mask, # rn,sn
'normal': normal, # rn,3
'mask': mask, # rn,1
}
return outputs
def render_with_loss(self, ray_batch, is_train, step):
render_outputs = self.render(ray_batch, is_train, step)
rgb_gt = ray_batch['rgb']
rgb_pr = render_outputs['rgb']
if self.rgb_loss == 'soft_l1':
epsilon = 0.001
rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
elif self.rgb_loss =='mse':
rgb_loss = F.mse_loss(rgb_pr, rgb_gt, reduction='none')
else:
raise NotImplementedError
rgb_loss = torch.mean(rgb_loss)
eikonal_loss = torch.sum(render_outputs['gradient_error'] * render_outputs['inner_mask']) / torch.sum(render_outputs['inner_mask'] + 1e-5)
loss = rgb_loss * self.lambda_rgb_loss + eikonal_loss * self.lambda_eikonal_loss
loss_batch = {
'eikonal': eikonal_loss,
'rendering': rgb_loss,
# 'mask': mask_loss,
}
if self.lambda_mask_loss>0 and self.use_mask:
mask_loss = F.mse_loss(render_outputs['mask'], ray_batch['mask'], reduction='none').mean()
loss += mask_loss * self.lambda_mask_loss
loss_batch['mask'] = mask_loss
return loss, loss_batch
class NeRFRenderer(BaseRenderer):
def __init__(self, train_batch_num, test_batch_num, bound=0.5, use_mask=False, lambda_rgb_loss=1.0, lambda_mask_loss=0.0):
super().__init__(train_batch_num, test_batch_num)
self.train_batch_num = train_batch_num
self.test_batch_num = test_batch_num
self.use_mask = use_mask
self.field = NGPNetwork(bound=bound)
self.update_interval = 16
self.fp16 = True
self.lambda_rgb_loss = lambda_rgb_loss
self.lambda_mask_loss = lambda_mask_loss
def render_impl(self, ray_batch, is_train, step):
rays_o, rays_d = ray_batch['rays_o'], ray_batch['rays_d']
with torch.cuda.amp.autocast(enabled=self.fp16):
if step % self.update_interval==0:
self.field.update_extra_state()
outputs = self.field.render(rays_o, rays_d,)
renderings={
'rgb': outputs['image'],
'depth': outputs['depth'],
'mask': outputs['weights_sum'].unsqueeze(-1),
}
return renderings
def render_with_loss(self, ray_batch, is_train, step):
render_outputs = self.render(ray_batch, is_train, step)
rgb_gt = ray_batch['rgb']
rgb_pr = render_outputs['rgb']
epsilon = 0.001
rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
rgb_loss = torch.mean(rgb_loss)
loss = rgb_loss * self.lambda_rgb_loss
loss_batch = {'rendering': rgb_loss}
if self.use_mask:
mask_loss = F.mse_loss(render_outputs['mask'], ray_batch['mask'], reduction='none')
mask_loss = torch.mean(mask_loss)
loss = loss + mask_loss * self.lambda_mask_loss
loss_batch['mask'] = mask_loss
return loss, loss_batch
def cartesian_to_spherical(xyz):
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
xy = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
z = np.sqrt(xy + xyz[:, 2] ** 2)
theta = np.arctan2(np.sqrt(xy), xyz[:, 2]) # for elevation angle defined from Z-axis down
# ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
azimuth = np.arctan2(xyz[:, 1], xyz[:, 0])
return np.array([theta, azimuth, z])
def get_pose(target_RT):
R, T = target_RT[:3, :3], target_RT[:, -1]
T_target = -R.T @ T
theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
return theta_target, azimuth_target, z_target
class RendererTrainer(pl.LightningModule):
def __init__(self, image_path, data_path, total_steps, warm_up_steps, log_dir, train_batch_fg_num=0,
use_cube_feats=False, cube_ckpt=None, cube_cfg=None, cube_bound=0.5,
train_batch_num=4096, test_batch_num=8192, use_warm_up=True, use_mask=True,
lambda_rgb_loss=1.0, lambda_mask_loss=0.0, renderer='neus',
# used in neus
lambda_eikonal_loss=0.1,
coarse_sn=64, fine_sn=64):
super().__init__()
self.num_images = 36 # todo ours 36, syncdreamer 16
self.image_size = 256
self.log_dir = log_dir
(Path(log_dir)/'images').mkdir(exist_ok=True, parents=True)
self.train_batch_num = train_batch_num
self.train_batch_fg_num = train_batch_fg_num
self.test_batch_num = test_batch_num
self.image_path = image_path
self.data_path = data_path
self.total_steps = total_steps
self.warm_up_steps = warm_up_steps
self.use_mask = use_mask
self.lambda_eikonal_loss = lambda_eikonal_loss
self.lambda_rgb_loss = lambda_rgb_loss
self.lambda_mask_loss = lambda_mask_loss
self.use_warm_up = use_warm_up
self.use_cube_feats, self.cube_cfg, self.cube_ckpt = use_cube_feats, cube_cfg, cube_ckpt
self._init_dataset()
if renderer=='neus':
self.renderer = NeuSRenderer(train_batch_num, test_batch_num,
lambda_rgb_loss=lambda_rgb_loss,
lambda_eikonal_loss=lambda_eikonal_loss,
lambda_mask_loss=lambda_mask_loss,
coarse_sn=coarse_sn, fine_sn=fine_sn)
elif renderer=='ngp':
self.renderer = NeRFRenderer(train_batch_num, test_batch_num, bound=cube_bound, use_mask=use_mask, lambda_mask_loss=lambda_mask_loss, lambda_rgb_loss=lambda_rgb_loss,)
else:
raise NotImplementedError
self.validation_index = 0
def _construct_ray_batch(self, images_info):
image_num = images_info['images'].shape[0]
_, h, w, _ = images_info['images'].shape
coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)] # h,w,2
coords = coords.float()[None, :, :, :].repeat(image_num, 1, 1, 1) # imn,h,w,2
coords = coords.reshape(image_num, h * w, 2)
coords = torch.cat([coords, torch.ones(image_num, h * w, 1, dtype=torch.float32)], 2) # imn,h*w,3
# imn,h*w,3 @ imn,3,3 => imn,h*w,3
rays_d = coords @ torch.inverse(images_info['Ks']).permute(0, 2, 1)
poses = images_info['poses'] # imn,3,4
R, t = poses[:, :, :3], poses[:, :, 3:]
rays_d = rays_d @ R
rays_d = F.normalize(rays_d, dim=-1)
rays_o = -R.permute(0,2,1) @ t # imn,3,3 @ imn,3,1
rays_o = rays_o.permute(0, 2, 1).repeat(1, h*w, 1) # imn,h*w,3
ray_batch = {
'rgb': images_info['images'].reshape(image_num*h*w,3),
'mask': images_info['masks'].reshape(image_num*h*w,1),
'rays_o': rays_o.reshape(image_num*h*w,3).float(),
'rays_d': rays_d.reshape(image_num*h*w,3).float(),
}
return ray_batch
@staticmethod
def load_model(cfg, ckpt):
config = OmegaConf.load(cfg)
model = instantiate_from_config(config.model)
print(f'loading model from {ckpt} ...')
ckpt = torch.load(ckpt)
model.load_state_dict(ckpt['state_dict'])
model = model.cuda().eval()
return model
def _init_dataset(self):
mask_predictor = BackgroundRemoval()
# syncdreamer fixed 16 views
# self.K, self.azs, self.els, self.dists, self.poses = read_pickle(f'meta_info/camera-{self.num_images}.pkl')
# for ours+NeuS, we pre fix 36 views
self.K = np.array([[280.,0.,128.],[0.,280.,128.],[0.,0.,1.]], dtype=np.float32)
data_dir = os.path.join(self.data_path, "mario/render_sync_36_single/model/") # fixed 36 views
# get all files .npy
self.azs = []
self.els = []
self.dists = []
self.poses = []
for index in range(self.num_images):
pose = np.load(os.path.join(data_dir, "%03d.npy"%index))[:3, :] # in blender
self.poses.append(pose)
theta, azimuth, radius = get_pose(pose)
self.azs.append(azimuth)
self.els.append(theta)
self.dists.append(radius)
# stack to numpy along axis 0
self.azs = np.stack(self.azs, axis=0) # [25,]
self.els = np.stack(self.els, axis=0) # [25,]
self.dists = np.stack(self.dists, axis=0) # [25,]
self.poses = np.stack(self.poses, axis=0) # [25, 3, 4]
self.images_info = {'images': [] ,'masks': [], 'Ks': [], 'poses':[]}
img = imread(self.image_path)
for index in range(self.num_images):
rgb = np.copy(img[:,index*self.image_size:(index+1)*self.image_size,:])
# predict mask
if self.use_mask:
imsave(f'{self.log_dir}/input-{index}.png', rgb)
masked_image = mask_predictor(rgb)
imsave(f'{self.log_dir}/masked-{index}.png', masked_image)
mask = masked_image[:,:,3].astype(np.float32)/255
else:
h, w, _ = rgb.shape
mask = np.zeros([h,w], np.float32)
rgb = rgb.astype(np.float32)/255
K, pose = np.copy(self.K), self.poses[index]
self.images_info['images'].append(torch.from_numpy(rgb.astype(np.float32))) # h,w,3
self.images_info['masks'].append(torch.from_numpy(mask.astype(np.float32))) # h,w
self.images_info['Ks'].append(torch.from_numpy(K.astype(np.float32)))
self.images_info['poses'].append(torch.from_numpy(pose.astype(np.float32)))
for k, v in self.images_info.items(): self.images_info[k] = torch.stack(v, 0) # stack all values
self.train_batch = self._construct_ray_batch(self.images_info)
self.train_batch_pseudo_fg = {}
pseudo_fg_mask = torch.sum(self.train_batch['rgb']>0.99,1)!=3
for k, v in self.train_batch.items():
self.train_batch_pseudo_fg[k] = v[pseudo_fg_mask]
self.train_ray_fg_num = int(torch.sum(pseudo_fg_mask).cpu().numpy())
self.train_ray_num = self.num_images * self.image_size ** 2
self._shuffle_train_batch()
self._shuffle_train_fg_batch()
def _shuffle_train_batch(self):
self.train_batch_i = 0
shuffle_idxs = torch.randperm(self.train_ray_num, device='cpu') # shuffle
for k, v in self.train_batch.items():
self.train_batch[k] = v[shuffle_idxs]
def _shuffle_train_fg_batch(self):
self.train_batch_fg_i = 0
shuffle_idxs = torch.randperm(self.train_ray_fg_num, device='cpu') # shuffle
for k, v in self.train_batch_pseudo_fg.items():
self.train_batch_pseudo_fg[k] = v[shuffle_idxs]
def training_step(self, batch, batch_idx):
train_ray_batch = {k: v[self.train_batch_i:self.train_batch_i + self.train_batch_num].cuda() for k, v in self.train_batch.items()}
self.train_batch_i += self.train_batch_num
if self.train_batch_i + self.train_batch_num >= self.train_ray_num: self._shuffle_train_batch()
if self.train_batch_fg_num>0:
train_ray_batch_fg = {k: v[self.train_batch_fg_i:self.train_batch_fg_i+self.train_batch_fg_num].cuda() for k, v in self.train_batch_pseudo_fg.items()}
self.train_batch_fg_i += self.train_batch_fg_num
if self.train_batch_fg_i + self.train_batch_fg_num >= self.train_ray_fg_num: self._shuffle_train_fg_batch()
for k, v in train_ray_batch_fg.items():
train_ray_batch[k] = torch.cat([train_ray_batch[k], v], 0)
loss, loss_batch = self.renderer.render_with_loss(train_ray_batch, is_train=True, step=self.global_step)
self.log_dict(loss_batch, prog_bar=True, logger=True, on_step=True, on_epoch=False, rank_zero_only=True)
self.log('step', self.global_step, prog_bar=True, on_step=True, on_epoch=False, logger=False, rank_zero_only=True)
lr = self.optimizers().param_groups[0]['lr']
self.log('lr', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False, rank_zero_only=True)
return loss
def _slice_images_info(self, index):
return {k:v[index:index+1] for k, v in self.images_info.items()}
@torch.no_grad()
def validation_step(self, batch, batch_idx):
with torch.no_grad():
if self.global_rank==0:
# we output an rendering image
images_info = self._slice_images_info(self.validation_index)
self.validation_index += 1
self.validation_index %= self.num_images
test_ray_batch = self._construct_ray_batch(images_info)
test_ray_batch = {k: v.cuda() for k,v in test_ray_batch.items()}
test_ray_batch['near'], test_ray_batch['far'] = near_far_from_sphere(test_ray_batch['rays_o'], test_ray_batch['rays_d'])
render_outputs = self.renderer.render(test_ray_batch, False, self.global_step)
process = lambda x: (x.cpu().numpy() * 255).astype(np.uint8)
h, w = self.image_size, self.image_size
rgb = torch.clamp(render_outputs['rgb'].reshape(h, w, 3), max=1.0, min=0.0)
mask = torch.clamp(render_outputs['mask'].reshape(h, w, 1), max=1.0, min=0.0)
mask_ = torch.repeat_interleave(mask, 3, dim=-1)
output_image = concat_images_list(process(rgb), process(mask_))
if 'normal' in render_outputs:
normal = torch.clamp((render_outputs['normal'].reshape(h, w, 3) + 1) / 2, max=1.0, min=0.0)
normal = normal * mask # we only show foregound normal
output_image = concat_images_list(output_image, process(normal))
# save images
imsave(f'{self.log_dir}/images/{self.global_step}.jpg', output_image)
def configure_optimizers(self):
lr = self.learning_rate
opt = torch.optim.AdamW([{"params": self.renderer.parameters(), "lr": lr},], lr=lr)
def schedule_fn(step):
total_step = self.total_steps
warm_up_step = self.warm_up_steps
warm_up_init = 0.02
warm_up_end = 1.0
final_lr = 0.02
interval = 1000
times = total_step // interval
ratio = np.power(final_lr, 1/times)
if step<warm_up_step:
learning_rate = (step / warm_up_step) * (warm_up_end - warm_up_init) + warm_up_init
else:
learning_rate = ratio ** (step // interval) * warm_up_end
return learning_rate
if self.use_warm_up:
scheduler = [{
'scheduler': LambdaLR(opt, lr_lambda=schedule_fn),
'interval': 'step',
'frequency': 1
}]
else:
scheduler = []
return [opt], scheduler
================================================
FILE: 3drecon/run_NeuS.py
================================================
import os
import numpy as np
from tqdm import tqdm
# ours + NeuS
DATA_DIR = "/home/xin/data/EscherNet/Data/GSO30" # GSO
exp_dir = "/home/xin/6DoF/GSO3D/"
config = "configs/neus_36.yaml"
exps = [1]
# exps = [1, 2, 3, 5, 10]
for exp in exps:
OUTPUT_DIR = os.path.join(exp_dir, f"logs_GSO_T{exp}M36_99k")
output_NeuS = f"ours_GSO_T{exp}"
os.makedirs(output_NeuS, exist_ok=True)
obj_names = os.listdir(DATA_DIR)
for obj_name in tqdm(obj_names):
if os.path.exists(os.path.join(output_NeuS, "NeuS", obj_name, "mesh.ply")):
print("NeuS already trained for: ", obj_name)
continue
# remove the folder for new training
os.system(f"rm -rf {output_NeuS}/NeuS/{obj_name}")
print("Training NeuS for: ", obj_name)
input_img = os.path.join(OUTPUT_DIR, obj_name, "0.png")
# input_img = os.path.join(OUTPUT_DIR, obj_name, "gt.png") # ground truth image
cmd = f"python train_renderer.py -i {input_img} \
-d {DATA_DIR} \
-n {obj_name} \
-b {config} \
-l {output_NeuS}/NeuS"
os.system(cmd)
================================================
FILE: 3drecon/train_renderer.py
================================================
import argparse
import imageio
import numpy as np
import torch
import torch.nn.functional as F
from pathlib import Path
import trimesh
from omegaconf import OmegaConf
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, Callback
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning import Trainer
from skimage.io import imsave
from tqdm import tqdm
import mcubes
from renderer.renderer import NeuSRenderer, DEFAULT_SIDE_LENGTH
from util import instantiate_from_config, read_pickle
class ResumeCallBacks(Callback):
def __init__(self):
pass
def on_train_start(self, trainer, pl_module):
pl_module.optimizers().param_groups = pl_module.optimizers()._optimizer.param_groups
def render_images(model, output,):
# render from model
n = 180
azimuths = (np.arange(n) / n * np.pi * 2).astype(np.float32)
elevations = np.deg2rad(np.asarray([30] * n).astype(np.float32))
K, _, _, _, poses = read_pickle(f'meta_info/camera-16.pkl')
output_points
h, w = 256, 256
default_size = 256
K = np.diag([w/default_size,h/default_size,1.0]) @ K
imgs = []
for ni in tqdm(range(n)):
# R = euler2mat(azimuths[ni], elevations[ni], 0, 'szyx')
# R = np.asarray([[0,-1,0],[0,0,-1],[1,0,0]]) @ R
e, a = elevations[ni], azimuths[ni]
row1 = np.asarray([np.sin(e)*np.cos(a),np.sin(e)*np.sin(a),-np.cos(e)])
row0 = np.asarray([-np.sin(a),np.cos(a), 0])
row2 = np.cross(row0, row1)
R = np.stack([row0,row1,row2],0)
t = np.asarray([0,0,1.5])
pose = np.concatenate([R,t[:,None]],1)
pose_ = torch.from_numpy(pose.astype(np.float32)).unsqueeze(0)
K_ = torch.from_numpy(K.astype(np.float32)).unsqueeze(0) # [1,3,3]
coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)] # h,w,2
coords = coords.float()[None, :, :, :].repeat(1, 1, 1, 1) # imn,h,w,2
coords = coords.reshape(1, h * w, 2)
coords = torch.cat([coords, torch.ones(1, h * w, 1, dtype=torch.float32)], 2) # imn,h*w,3
# imn,h*w,3 @ imn,3,3 => imn,h*w,3
rays_d = coords @ torch.inverse(K_).permute(0, 2, 1)
R, t = pose_[:, :, :3], pose_[:, :, 3:]
rays_d = rays_d @ R
rays_d = F.normalize(rays_d, dim=-1)
rays_o = -R.permute(0, 2, 1) @ t # imn,3,3 @ imn,3,1
rays_o = rays_o.permute(0, 2, 1).repeat(1, h * w, 1) # imn,h*w,3
ray_batch = {
'rays_o': rays_o.reshape(-1,3).cuda(),
'rays_d': rays_d.reshape(-1,3).cuda(),
}
with torch.no_grad():
image = model.renderer.render(ray_batch,False,5000)['rgb'].reshape(h,w,3)
image = (image.cpu().numpy() * 255).astype(np.uint8)
imgs.append(image)
imageio.mimsave(f'{output}/rendering.mp4', imgs, fps=30)
def extract_fields(bound_min, bound_max, resolution, query_func, batch_size=64, outside_val=1.0):
N = batch_size
X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
with torch.no_grad():
for xi, xs in enumerate(X):
for yi, ys in enumerate(Y):
for zi, zs in enumerate(Z):
xx, yy, zz = torch.meshgrid(xs, ys, zs)
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1).cuda()
val = query_func(pts).detach()
outside_mask = torch.norm(pts,dim=-1)>=1.0
val[outside_mask]=outside_val
val = val.reshape(len(xs), len(ys), len(zs)).cpu().numpy()
u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
return u
def extract_geometry(bound_min, bound_max, resolution, threshold, query_func, color_func, outside_val=1.0):
u = extract_fields(bound_min, bound_max, resolution, query_func, outside_val=outside_val)
vertices, triangles = mcubes.marching_cubes(u, threshold)
b_max_np = bound_max.detach().cpu().numpy()
b_min_np = bound_min.detach().cpu().numpy()
vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
vertex_colors = color_func(vertices)
return vertices, triangles, vertex_colors
def extract_mesh(model, output, resolution=512):
if not isinstance(model.renderer, NeuSRenderer): return
bbox_min = -torch.ones(3)*DEFAULT_SIDE_LENGTH
bbox_max = torch.ones(3)*DEFAULT_SIDE_LENGTH
with torch.no_grad():
vertices, triangles, vertex_colors = extract_geometry(bbox_min, bbox_max, resolution, 0, lambda x: model.renderer.sdf_network.sdf(x), lambda x: model.renderer.get_vertex_colors(x))
# output geometry
mesh = trimesh.Trimesh(vertices, triangles, vertex_colors=vertex_colors)
mesh.export(str(f'{output}/mesh.ply'))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--image_path', type=str, required=True)
parser.add_argument('-n', '--name', type=str, required=True)
parser.add_argument('-b', '--base', type=str, default='configs/neus.yaml')
parser.add_argument('-d', '--data_path', type=str, default='/data/GSO/')
parser.add_argument('-l', '--log', type=str, default='output/renderer')
parser.add_argument('-s', '--seed', type=int, default=6033)
parser.add_argument('-g', '--gpus', type=str, default='0,')
parser.add_argument('-r', '--resume', action='store_true', default=False, dest='resume')
parser.add_argument('--fp16', action='store_true', default=False, dest='fp16')
opt = parser.parse_args()
# seed_everything(opt.seed)
# configs
cfg = OmegaConf.load(opt.base)
name = opt.name
log_dir, ckpt_dir = Path(opt.log) / name, Path(opt.log) / name / 'ckpt'
cfg.model.params['image_path'] = opt.image_path
cfg.model.params['log_dir'] = log_dir
cfg.model.params['data_path'] = opt.data_path
# setup
log_dir.mkdir(exist_ok=True, parents=True)
ckpt_dir.mkdir(exist_ok=True, parents=True)
trainer_config = cfg.trainer
callback_config = cfg.callbacks
model_config = cfg.model
data_config = cfg.data
data_config.params.seed = opt.seed
data = instantiate_from_config(data_config)
data.prepare_data()
data.setup('fit')
model = instantiate_from_config(model_config,)
model.cpu()
model.learning_rate = model_config.base_lr
# logger
logger = TensorBoardLogger(save_dir=log_dir, name='tensorboard_logs')
callbacks=[]
callbacks.append(LearningRateMonitor(logging_interval='step'))
callbacks.append(ModelCheckpoint(dirpath=ckpt_dir, filename="{epoch:06}", verbose=True, save_last=True, every_n_train_steps=callback_config.save_interval))
# trainer
trainer_config.update({
"accelerator": "cuda", "check_val_every_n_epoch": None,
"benchmark": True, "num_sanity_val_steps": 0,
"devices": 1, "gpus": opt.gpus,
})
if opt.fp16:
trainer_config['precision']=16
if opt.resume:
callbacks.append(ResumeCallBacks())
trainer_config['resume_from_checkpoint'] = str(ckpt_dir / 'last.ckpt')
else:
if (ckpt_dir / 'last.ckpt').exists():
raise RuntimeError(f"checkpoint {ckpt_dir / 'last.ckpt'} existing ...")
trainer = Trainer.from_argparse_args(args=argparse.Namespace(), **trainer_config, logger=logger, callbacks=callbacks)
trainer.fit(model, data)
model = model.cuda().eval()
# render_images(model, log_dir)
extract_mesh(model, log_dir)
if __name__=="__main__":
main()
================================================
FILE: 3drecon/util.py
================================================
import importlib
import pickle
import numpy as np
import cv2
def instantiate_from_config(config):
if not "target" in config:
if config == '__is_first_stage__':
return None
elif config == "__is_unconditional__":
return None
raise KeyError("Expected key `target` to instantiate.")
return get_obj_from_str(config["target"])(**config.get("params", dict()))
def get_obj_from_str(string, reload=False):
module, cls = string.rsplit(".", 1)
if reload:
module_imp = importlib.import_module(module)
importlib.reload(module_imp)
return getattr(importlib.import_module(module, package=None), cls)
def read_pickle(pkl_path):
with open(pkl_path, 'rb') as f:
return pickle.load(f)
def output_points(fn,pts,colors=None):
with open(fn, 'w') as f:
for pi, pt in enumerate(pts):
f.write(f'{pt[0]:.6f} {pt[1]:.6f} {pt[2]:.6f} ')
if colors is not None:
f.write(f'{int(colors[pi,0])} {int(colors[pi,1])} {int(colors[pi,2])}')
f.write('\n')
def concat_images(img0,img1,vert=False):
if not vert:
h0,h1=img0.shape[0],img1.shape[0],
if h0<h1: img0=cv2.copyMakeBorder(img0,0,h1-h0,0,0,borderType=cv2.BORDER_CONSTANT,value=0)
if h1<h0: img1=cv2.copyMakeBorder(img1,0,h0-h1,0,0,borderType=cv2.BORDER_CONSTANT,value=0)
img = np.concatenate([img0, img1], axis=1)
else:
w0,w1=img0.shape[1],img1.shape[1]
if w0<w1: img0=cv2.copyMakeBorder(img0,0,0,0,w1-w0,borderType=cv2.BORDER_CONSTANT,value=0)
if w1<w0: img1=cv2.copyMakeBorder(img1,0,0,0,w0-w1,borderType=cv2.BORDER_CONSTANT,value=0)
img = np.concatenate([img0, img1], axis=0)
return img
def concat_images_list(*args,vert=False):
if len(args)==1: return args[0]
img_out=args[0]
for img in args[1:]:
img_out=concat_images(img_out,img,vert)
return img_out
================================================
FILE: 4DoF/CN_encoder.py
================================================
from transformers import ConvNextV2Model
import torch
from typing import Optional
import einops
class CN_encoder(ConvNextV2Model):
def __init__(self, config):
super().__init__(config)
def forward(
self,
pixel_values: torch.FloatTensor = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
):
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if pixel_values is None:
raise ValueError("You have to specify pixel_values")
embedding_output = self.embeddings(pixel_values)
encoder_outputs = self.encoder(
embedding_output,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
last_hidden_state = encoder_outputs[0]
image_embeddings = einops.rearrange(last_hidden_state, 'b c h w -> b (h w) c')
image_embeddings = self.layernorm(image_embeddings)
return image_embeddings
================================================
FILE: 4DoF/dataset.py
================================================
import os
import math
from pathlib import Path
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
import webdataset as wds
from torch.utils.data.distributed import DistributedSampler
import matplotlib.pyplot as plt
import sys
class ObjaverseDataLoader():
def __init__(self, root_dir, batch_size, total_view=12, num_workers=4):
self.root_dir = root_dir
self.batch_size = batch_size
self.num_workers = num_workers
self.total_view = total_view
image_transforms = [torchvision.transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])]
self.image_transforms = torchvision.transforms.Compose(image_transforms)
def train_dataloader(self):
dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=False,
image_transforms=self.image_transforms)
# sampler = DistributedSampler(dataset)
return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
# sampler=sampler)
def val_dataloader(self):
dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=True,
image_transforms=self.image_transforms)
sampler = DistributedSampler(dataset)
return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
def cartesian_to_spherical(xyz):
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
xy = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
z = np.sqrt(xy + xyz[:, 2] ** 2)
theta = np.arctan2(np.sqrt(xy), xyz[:, 2]) # for elevation angle defined from Z-axis down
# ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
azimuth = np.arctan2(xyz[:, 1], xyz[:, 0])
return np.array([theta, azimuth, z])
def get_pose(target_RT):
target_RT = target_RT[:3, :]
R, T = target_RT[:3, :3], target_RT[:, -1]
T_target = -R.T @ T
theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
# assert if z_target is out of range
if z_target.item() < 1.5 or z_target.item() > 2.2:
# print('z_target out of range 1.5-2.2', z_target.item())
z_target = np.clip(z_target.item(), 1.5, 2.2)
# with log scale for radius
target_T = torch.tensor([theta_target.item(), azimuth_target.item(), (np.log(z_target.item()) - np.log(1.5))/(np.log(2.2)-np.log(1.5)) * torch.pi, torch.tensor(0)])
assert torch.all(target_T <= torch.pi) and torch.all(target_T >= -torch.pi)
return target_T.numpy()
class ObjaverseData(Dataset):
def __init__(self,
root_dir='.objaverse/hf-objaverse-v1/views',
image_transforms=None,
total_view=12,
validation=False,
T_in=1,
T_out=1,
fix_sample=False,
) -> None:
"""Create a dataset from a folder of images.
If you pass in a root directory it will be searched for images
ending in ext (ext can be a list)
"""
self.root_dir = Path(root_dir)
self.total_view = total_view
self.T_in = T_in
self.T_out = T_out
self.fix_sample = fix_sample
self.paths = []
# # include all folders
# for folder in os.listdir(self.root_dir):
# if os.path.isdir(os.path.join(self.root_dir, folder)):
# self.paths.append(folder)
# load ids from .npy so we have exactly the same ids/order
self.paths = np.load("../scripts/obj_ids.npy")
# # only use 100K objects for ablation study
# self.paths = self.paths[:100000]
total_objects = len(self.paths)
assert total_objects == 790152, 'total objects %d' % total_objects
if validation:
self.paths = self.paths[math.floor(total_objects / 100. * 99.):] # used last 1% as validation
else:
self.paths = self.paths[:math.floor(total_objects / 100. * 99.)] # used first 99% as training
print('============= length of dataset %d =============' % len(self.paths))
self.tform = image_transforms
downscale = 512 / 256.
self.fx = 560. / downscale
self.fy = 560. / downscale
self.intrinsic = torch.tensor([[self.fx, 0, 128., 0, self.fy, 128., 0, 0, 1.]], dtype=torch.float64).view(3, 3)
def __len__(self):
return len(self.paths)
def cartesian_to_spherical(self, xyz):
ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
xy = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
z = np.sqrt(xy + xyz[:, 2] ** 2)
theta = np.arctan2(np.sqrt(xy), xyz[:, 2]) # for elevation angle defined from Z-axis down
# ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
azimuth = np.arctan2(xyz[:, 1], xyz[:, 0])
return np.array([theta, azimuth, z])
def get_T(self, target_RT, cond_RT):
R, T = target_RT[:3, :3], target_RT[:, -1]
T_target = -R.T @ T
R, T = cond_RT[:3, :3], cond_RT[:, -1]
T_cond = -R.T @ T
theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
d_theta = theta_target - theta_cond
d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
d_z = z_target - z_cond
d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
return d_T
def get_pose(self, target_RT):
R, T = target_RT[:3, :3], target_RT[:, -1]
T_target = -R.T @ T
theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
# assert if z_target is out of range
if z_target.item() < 1.5 or z_target.item() > 2.2:
# print('z_target out of range 1.5-2.2', z_target.item())
z_target = np.clip(z_target.item(), 1.5, 2.2)
# with log scale for radius
target_T = torch.tensor([theta_target.item(), azimuth_target.item(), (np.log(z_target.item()) - np.log(1.5))/(np.log(2.2)-np.log(1.5)) * torch.pi, torch.tensor(0)])
assert torch.all(target_T <= torch.pi) and torch.all(target_T >= -torch.pi)
return target_T
def load_im(self, path, color):
'''
replace background pixel with random color in rendering
'''
try:
img = plt.imread(path)
except:
print(path)
sys.exit()
img[img[:, :, -1] == 0.] = color
img = Image.fromarray(np.uint8(img[:, :, :3] * 255.))
return img
def __getitem__(self, index):
data = {}
total_view = 12
if self.fix_sample:
if self.T_out > 1:
indexes = range(total_view)
index_targets = list(indexes[:2]) + list(indexes[-(self.T_out-2):])
index_inputs = indexes[1:self.T_in+1] # one overlap identity
else:
indexes = range(total_view)
index_targets = indexes[:self.T_out]
index_inputs = indexes[self.T_out-1:self.T_in+self.T_out-1] # one overlap identity
else:
assert self.T_in + self.T_out <= total_view
# training with replace, including identity
indexes = np.random.choice(range(total_view), self.T_in+self.T_out, replace=True)
index_inputs = indexes[:self.T_in]
index_targets = indexes[self.T_in:]
filename = os.path.join(self.root_dir, self.paths[index])
color = [1., 1., 1., 1.]
try:
input_ims = []
target_ims = []
target_Ts = []
cond_Ts = []
for i, index_input in enumerate(index_inputs):
input_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_input), color))
input_ims.append(input_im)
input_RT = np.load(os.path.join(filename, '%03d.npy' % index_input))
cond_Ts.append(self.get_pose(input_RT))
for i, index_target in enumerate(index_targets):
target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
target_ims.append(target_im)
target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
target_Ts.append(self.get_pose(target_RT))
except:
print('error loading data ', filename)
filename = os.path.join(self.root_dir, '0a01f314e2864711aa7e33bace4bd8c8') # this one we know is valid
input_ims = []
target_ims = []
target_Ts = []
cond_Ts = []
# very hacky solution, sorry about this
for i, index_input in enumerate(index_inputs):
input_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_input), color))
input_ims.append(input_im)
input_RT = np.load(os.path.join(filename, '%03d.npy' % index_input))
cond_Ts.append(self.get_pose(input_RT))
for i, index_target in enumerate(index_targets):
target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
target_ims.append(target_im)
target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
target_Ts.append(self.get_pose(target_RT))
# stack to batch
data['image_input'] = torch.stack(input_ims, dim=0)
data['image_target'] = torch.stack(target_ims, dim=0)
data['pose_out'] = torch.stack(target_Ts, dim=0)
data['pose_in'] = torch.stack(cond_Ts, dim=0)
return data
def process_im(self, im):
im = im.convert("RGB")
return self.tform(im)
================================================
FILE: 4DoF/diffusers/__init__.py
================================================
__version__ = "0.18.2"
from .configuration_utils import ConfigMixin
from .utils import (
OptionalDependencyNotAvailable,
is_flax_available,
is_inflect_available,
is_invisible_watermark_available,
is_k_diffusion_available,
is_k_diffusion_version,
is_librosa_available,
is_note_seq_available,
is_onnx_available,
is_scipy_available,
is_torch_available,
is_torchsde_available,
is_transformers_available,
is_transformers_version,
is_unidecode_available,
logging,
)
try:
if not is_onnx_available():
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_onnx_objects import * # noqa F403
else:
from .pipelines import OnnxRuntimeModel
try:
if not is_torch_available():
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_pt_objects import * # noqa F403
else:
from .models import (
AutoencoderKL,
ControlNetModel,
ModelMixin,
PriorTransformer,
T5FilmDecoder,
Transformer2DModel,
UNet1DModel,
UNet2DConditionModel,
UNet2DModel,
UNet3DConditionModel,
VQModel,
)
from .optimization import (
get_constant_schedule,
get_constant_schedule_with_warmup,
get_cosine_schedule_with_warmup,
get_cosine_with_hard_restarts_schedule_with_warmup,
get_linear_schedule_with_warmup,
get_polynomial_decay_schedule_with_warmup,
get_scheduler,
)
from .pipelines import (
AudioPipelineOutput,
ConsistencyModelPipeline,
DanceDiffusionPipeline,
DDIMPipeline,
DDPMPipeline,
DiffusionPipeline,
DiTPipeline,
ImagePipelineOutput,
KarrasVePipeline,
LDMPipeline,
LDMSuperResolutionPipeline,
PNDMPipeline,
RePaintPipeline,
ScoreSdeVePipeline,
)
from .schedulers import (
CMStochasticIterativeScheduler,
DDIMInverseScheduler,
DDIMParallelScheduler,
DDIMScheduler,
DDPMParallelScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepInverseScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
IPNDMScheduler,
KarrasVeScheduler,
KDPM2AncestralDiscreteScheduler,
KDPM2DiscreteScheduler,
PNDMScheduler,
RePaintScheduler,
SchedulerMixin,
ScoreSdeVeScheduler,
UnCLIPScheduler,
UniPCMultistepScheduler,
VQDiffusionScheduler,
)
from .training_utils import EMAModel
try:
if not (is_torch_available() and is_scipy_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_scipy_objects import * # noqa F403
else:
from .schedulers import LMSDiscreteScheduler
try:
if not (is_torch_available() and is_torchsde_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_torchsde_objects import * # noqa F403
else:
from .schedulers import DPMSolverSDEScheduler
try:
if not (is_torch_available() and is_transformers_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
from .pipelines import (
AltDiffusionImg2ImgPipeline,
AltDiffusionPipeline,
AudioLDMPipeline,
CycleDiffusionPipeline,
IFImg2ImgPipeline,
IFImg2ImgSuperResolutionPipeline,
IFInpaintingPipeline,
IFInpaintingSuperResolutionPipeline,
IFPipeline,
IFSuperResolutionPipeline,
ImageTextPipelineOutput,
KandinskyImg2ImgPipeline,
KandinskyInpaintPipeline,
KandinskyPipeline,
KandinskyPriorPipeline,
KandinskyV22ControlnetImg2ImgPipeline,
KandinskyV22ControlnetPipeline,
KandinskyV22Img2ImgPipeline,
KandinskyV22InpaintPipeline,
KandinskyV22Pipeline,
KandinskyV22PriorEmb2EmbPipeline,
KandinskyV22PriorPipeline,
LDMTextToImagePipeline,
PaintByExamplePipeline,
SemanticStableDiffusionPipeline,
ShapEImg2ImgPipeline,
ShapEPipeline,
StableDiffusionAttendAndExcitePipeline,
StableDiffusionControlNetImg2ImgPipeline,
StableDiffusionControlNetInpaintPipeline,
StableDiffusionControlNetPipeline,
StableDiffusionDepth2ImgPipeline,
StableDiffusionDiffEditPipeline,
StableDiffusionImageVariationPipeline,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipeline,
StableDiffusionInpaintPipelineLegacy,
StableDiffusionInstructPix2PixPipeline,
StableDiffusionLatentUpscalePipeline,
StableDiffusionLDM3DPipeline,
StableDiffusionModelEditingPipeline,
StableDiffusionPanoramaPipeline,
StableDiffusionParadigmsPipeline,
StableDiffusionPipeline,
StableDiffusionPipelineSafe,
StableDiffusionPix2PixZeroPipeline,
StableDiffusionSAGPipeline,
StableDiffusionUpscalePipeline,
StableUnCLIPImg2ImgPipeline,
StableUnCLIPPipeline,
TextToVideoSDPipeline,
TextToVideoZeroPipeline,
UnCLIPImageVariationPipeline,
UnCLIPPipeline,
UniDiffuserModel,
UniDiffuserPipeline,
UniDiffuserTextDecoder,
VersatileDiffusionDualGuidedPipeline,
VersatileDiffusionImageVariationPipeline,
VersatileDiffusionPipeline,
VersatileDiffusionTextToImagePipeline,
VideoToVideoSDPipeline,
VQDiffusionPipeline,
)
try:
if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import * # noqa F403
else:
from .pipelines import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline
try:
if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import * # noqa F403
else:
from .pipelines import StableDiffusionKDiffusionPipeline
try:
if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403
else:
from .pipelines import (
OnnxStableDiffusionImg2ImgPipeline,
OnnxStableDiffusionInpaintPipeline,
OnnxStableDiffusionInpaintPipelineLegacy,
OnnxStableDiffusionPipeline,
OnnxStableDiffusionUpscalePipeline,
StableDiffusionOnnxPipeline,
)
try:
if not (is_torch_available() and is_librosa_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_torch_and_librosa_objects import * # noqa F403
else:
from .pipelines import AudioDiffusionPipeline, Mel
try:
if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403
else:
from .pipelines import SpectrogramDiffusionPipeline
try:
if not is_flax_available():
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_flax_objects import * # noqa F403
else:
from .models.controlnet_flax import FlaxControlNetModel
from .models.modeling_flax_utils import FlaxModelMixin
from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
from .models.vae_flax import FlaxAutoencoderKL
from .pipelines import FlaxDiffusionPipeline
from .schedulers import (
FlaxDDIMScheduler,
FlaxDDPMScheduler,
FlaxDPMSolverMultistepScheduler,
FlaxKarrasVeScheduler,
FlaxLMSDiscreteScheduler,
FlaxPNDMScheduler,
FlaxSchedulerMixin,
FlaxScoreSdeVeScheduler,
)
try:
if not (is_flax_available() and is_transformers_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_flax_and_transformers_objects import * # noqa F403
else:
from .pipelines import (
FlaxStableDiffusionControlNetPipeline,
FlaxStableDiffusionImg2ImgPipeline,
FlaxStableDiffusionInpaintPipeline,
FlaxStableDiffusionPipeline,
)
try:
if not (is_note_seq_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_note_seq_objects import * # noqa F403
else:
from .pipelines import MidiProcessor
================================================
FILE: 4DoF/diffusers/commands/__init__.py
================================================
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import ABC, abstractmethod
from argparse import ArgumentParser
class BaseDiffusersCLICommand(ABC):
@staticmethod
@abstractmethod
def register_subcommand(parser: ArgumentParser):
raise NotImplementedError()
@abstractmethod
def run(self):
raise NotImplementedError()
================================================
FILE: 4DoF/diffusers/commands/diffusers_cli.py
================================================
#!/usr/bin/env python
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from argparse import ArgumentParser
from .env import EnvironmentCommand
def main():
parser = ArgumentParser("Diffusers CLI tool", usage="diffusers-cli <command> [<args>]")
commands_parser = parser.add_subparsers(help="diffusers-cli command helpers")
# Register commands
EnvironmentCommand.register_subcommand(commands_parser)
# Let's go
args = parser.parse_args()
if not hasattr(args, "func"):
parser.print_help()
exit(1)
# Run
service = args.func(args)
service.run()
if __name__ == "__main__":
main()
================================================
FILE: 4DoF/diffusers/commands/env.py
================================================
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import platform
from argparse import ArgumentParser
import huggingface_hub
from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand
def info_command_factory(_):
return EnvironmentCommand()
class EnvironmentCommand(BaseDiffusersCLICommand):
@staticmethod
def register_subcommand(parser: ArgumentParser):
download_parser = parser.add_parser("env")
download_parser.set_defaults(func=info_command_factory)
def run(self):
hub_version = huggingface_hub.__version__
pt_version = "not installed"
pt_cuda_available = "NA"
if is_torch_available():
import torch
pt_version = torch.__version__
pt_cuda_available = torch.cuda.is_available()
transformers_version = "not installed"
if is_transformers_available():
import transformers
transformers_version = transformers.__version__
accelerate_version = "not installed"
if is_accelerate_available():
import accelerate
accelerate_version = accelerate.__version__
xformers_version = "not installed"
if is_xformers_available():
import xformers
xformers_version = xformers.__version__
info = {
"`diffusers` version": version,
"Platform": platform.platform(),
"Python version": platform.python_version(),
"PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
"Huggingface_hub version": hub_version,
"Transformers version": transformers_version,
"Accelerate version": accelerate_version,
"xFormers version": xformers_version,
"Using GPU in script?": "<fill in>",
"Using distributed or parallel set-up in script?": "<fill in>",
}
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
print(self.format_dict(info))
return info
@staticmethod
def format_dict(d):
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
================================================
FILE: 4DoF/diffusers/configuration_utils.py
================================================
# coding=utf-8
# Copyright 2023 The HuggingFace Inc. team.
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" ConfigMixin base class and utilities."""
import dataclasses
import functools
import importlib
import inspect
import json
import os
import re
from collections import OrderedDict
from pathlib import PosixPath
from typing import Any, Dict, Tuple, Union
import numpy as np
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
from requests import HTTPError
from . import __version__
from .utils import (
DIFFUSERS_CACHE,
HUGGINGFACE_CO_RESOLVE_ENDPOINT,
DummyObject,
deprecate,
extract_commit_hash,
http_user_agent,
logging,
)
logger = logging.get_logger(__name__)
_re_configuration_file = re.compile(r"config\.(.*)\.json")
class FrozenDict(OrderedDict):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
for key, value in self.items():
setattr(self, key, value)
self.__frozen = True
def __delitem__(self, *args, **kwargs):
raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
def setdefault(self, *args, **kwargs):
raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
def pop(self, *args, **kwargs):
raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
def update(self, *args, **kwargs):
raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
def __setattr__(self, name, value):
if hasattr(self, "__frozen") and self.__frozen:
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
super().__setattr__(name, value)
def __setitem__(self, name, value):
if hasattr(self, "__frozen") and self.__frozen:
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
super().__setitem__(name, value)
class ConfigMixin:
r"""
Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
saving classes that inherit from [`ConfigMixin`].
Class attributes:
- **config_name** (`str`) -- A filename under which the config should stored when calling
[`~ConfigMixin.save_config`] (should be overridden by parent class).
- **ignore_for_config** (`List[str]`) -- A list of attributes that should not be saved in the config (should be
overridden by subclass).
- **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
- **_deprecated_kwargs** (`List[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
subclass).
"""
config_name = None
ignore_for_config = []
has_compatibles = False
_deprecated_kwargs = []
def register_to_config(self, **kwargs):
if self.config_name is None:
raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
# Special case for `kwargs` used in deprecation warning added to schedulers
# TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
# or solve in a more general way.
kwargs.pop("kwargs", None)
if not hasattr(self, "_internal_dict"):
internal_dict = kwargs
else:
previous_dict = dict(self._internal_dict)
internal_dict = {**self._internal_dict, **kwargs}
logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
self._internal_dict = FrozenDict(internal_dict)
def __getattr__(self, name: str) -> Any:
"""The only reason we overwrite `getattr` here is to gracefully deprecate accessing
config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
Tihs funtion is mostly copied from PyTorch's __getattr__ overwrite:
https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
"""
is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
is_attribute = name in self.__dict__
if is_in_config and not is_attribute:
deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'."
deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
return self._internal_dict[name]
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
"""
Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
[`~ConfigMixin.from_config`] class method.
Args:
save_directory (`str` or `os.PathLike`):
Directory where the configuration JSON file is saved (will be created if it does not exist).
"""
if os.path.isfile(save_directory):
raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
os.makedirs(save_directory, exist_ok=True)
# If we save using the predefined names, we can load using `from_config`
output_config_file = os.path.join(save_directory, self.config_name)
self.to_json_file(output_config_file)
logger.info(f"Configuration saved in {output_config_file}")
@classmethod
def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
r"""
Instantiate a Python class from a config dictionary.
Parameters:
config (`Dict[str, Any]`):
A config dictionary from which the Python class is instantiated. Make sure to only load configuration
files of compatible classes.
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
Whether kwargs that are not consumed by the Python class should be returned or not.
kwargs (remaining dictionary of keyword arguments, *optional*):
Can be used to update the configuration object (after it is loaded) and initiate the Python class.
`**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
overwrite the same named arguments in `config`.
Returns:
[`ModelMixin`] or [`SchedulerMixin`]:
A model or scheduler object instantiated from a config dictionary.
Examples:
```python
>>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
>>> # Download scheduler from huggingface.co and cache.
>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
>>> # Instantiate DDIM scheduler class with same config as DDPM
>>> scheduler = DDIMScheduler.from_config(scheduler.config)
>>> # Instantiate PNDM scheduler class with same config as DDPM
>>> scheduler = PNDMScheduler.from_config(scheduler.config)
```
"""
# <===== TO BE REMOVED WITH DEPRECATION
# TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
if "pretrained_model_name_or_path" in kwargs:
config = kwargs.pop("pretrained_model_name_or_path")
if config is None:
raise ValueError("Please make sure to provide a config as the first positional argument.")
# ======>
if not isinstance(config, dict):
deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
if "Scheduler" in cls.__name__:
deprecation_message += (
f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
" Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
" be removed in v1.0.0."
)
elif "Model" in cls.__name__:
deprecation_message += (
f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
" instead. This functionality will be removed in v1.0.0."
)
deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
# Allow dtype to be specified on initialization
if "dtype" in unused_kwargs:
init_dict["dtype"] = unused_kwargs.pop("dtype")
# add possible deprecated kwargs
for deprecated_kwarg in cls._deprecated_kwargs:
if deprecated_kwarg in unused_kwargs:
init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
# Return model and optionally state and/or unused_kwargs
model = cls(**init_dict)
# make sure to also save config parameters that might be used for compatible classes
model.register_to_config(**hidden_dict)
# add hidden kwargs of compatible classes to unused_kwargs
unused_kwargs = {**unused_kwargs, **hidden_dict}
if return_unused_kwargs:
return (model, unused_kwargs)
else:
return model
@classmethod
def get_config_dict(cls, *args, **kwargs):
deprecation_message = (
f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
" removed in version v1.0.0"
)
deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
return cls.load_config(*args, **kwargs)
@classmethod
def load_config(
cls,
pretrained_model_name_or_path: Union[str, os.PathLike],
return_unused_kwargs=False,
return_commit_hash=False,
**kwargs,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
r"""
Load a model or scheduler configuration.
Parameters:
pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
Can be either:
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
the Hub.
- A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
[`~ConfigMixin.save_config`].
cache_dir (`Union[str, os.PathLike]`, *optional*):
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
is not used.
force_download (`bool`, *optional*, defaults to `False`):
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
cached versions if they exist.
resume_download (`bool`, *optional*, defaults to `False`):
Whether or not to resume downloading the model weights and configuration files. If set to `False`, any
incompletely downloaded files are deleted.
proxies (`Dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
output_loading_info(`bool`, *optional*, defaults to `False`):
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
local_files_only (`bool`, *optional*, defaults to `False`):
Whether to only load local model weights and configuration files or not. If set to `True`, the model
won't be downloaded from the Hub.
use_auth_token (`str` or *bool*, *optional*):
The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
`diffusers-cli login` (stored in `~/.huggingface`) is used.
revision (`str`, *optional*, defaults to `"main"`):
The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
allowed by Git.
subfolder (`str`, *optional*, defaults to `""`):
The subfolder location of a model file within a larger model repository on the Hub or locally.
return_unused_kwargs (`bool`, *optional*, defaults to `False):
Whether unused keyword arguments of the config are returned.
return_commit_hash (`bool`, *optional*, defaults to `False):
Whether the `commit_hash` of the loaded configuration are returned.
Returns:
`dict`:
A dictionary of all the parameters stored in a JSON configuration file.
"""
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
proxies = kwargs.pop("proxies", None)
use_auth_token = kwargs.pop("use_auth_token", None)
local_files_only = kwargs.pop("local_files_only", False)
revision = kwargs.pop("revision", None)
_ = kwargs.pop("mirror", None)
subfolder = kwargs.pop("subfolder", None)
user_agent = kwargs.pop("user_agent", {})
user_agent = {**user_agent, "file_type": "config"}
user_agent = http_user_agent(user_agent)
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
if cls.config_name is None:
raise ValueError(
"`self.config_name` is not defined. Note that one should not load a config from "
"`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
)
if os.path.isfile(pretrained_model_name_or_path):
config_file = pretrained_model_name_or_path
elif os.path.isdir(pretrained_model_name_or_path):
if os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
# Load from a PyTorch checkpoint
config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
elif subfolder is not None and os.path.isfile(
os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
):
config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
else:
raise EnvironmentError(
f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
)
else:
try:
# Load from URL or cache if already cached
config_file = hf_hub_download(
pretrained_model_name_or_path,
filename=cls.config_name,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
local_files_only=local_files_only,
use_auth_token=use_auth_token,
user_agent=user_agent,
subfolder=subfolder,
revision=revision,
)
except RepositoryNotFoundError:
raise EnvironmentError(
f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
" listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
" token having permission to this repo with `use_auth_token`
gitextract_grlqwy99/
├── .gitignore
├── 3drecon/
│ ├── configs/
│ │ └── neus_36.yaml
│ ├── ours_GSO_T1/
│ │ └── NeuS/
│ │ ├── grandmother/
│ │ │ ├── mesh.ply
│ │ │ └── tensorboard_logs/
│ │ │ └── version_0/
│ │ │ ├── events.out.tfevents.1706660493.xin-PC.934205.0
│ │ │ └── hparams.yaml
│ │ ├── lion/
│ │ │ ├── mesh.ply
│ │ │ └── tensorboard_logs/
│ │ │ └── version_0/
│ │ │ ├── events.out.tfevents.1706659999.xin-PC.931603.0
│ │ │ └── hparams.yaml
│ │ └── train/
│ │ └── tensorboard_logs/
│ │ └── version_0/
│ │ ├── events.out.tfevents.1706660976.xin-PC.936418.0
│ │ └── hparams.yaml
│ ├── raymarching/
│ │ ├── __init__.py
│ │ ├── backend.py
│ │ ├── raymarching.py
│ │ ├── setup.py
│ │ └── src/
│ │ ├── bindings.cpp
│ │ ├── raymarching.cu
│ │ └── raymarching.h
│ ├── renderer/
│ │ ├── agg_net.py
│ │ ├── cost_reg_net.py
│ │ ├── dummy_dataset.py
│ │ ├── feature_net.py
│ │ ├── neus_networks.py
│ │ ├── ngp_renderer.py
│ │ └── renderer.py
│ ├── run_NeuS.py
│ ├── train_renderer.py
│ └── util.py
├── 4DoF/
│ ├── CN_encoder.py
│ ├── dataset.py
│ ├── diffusers/
│ │ ├── __init__.py
│ │ ├── commands/
│ │ │ ├── __init__.py
│ │ │ ├── diffusers_cli.py
│ │ │ └── env.py
│ │ ├── configuration_utils.py
│ │ ├── dependency_versions_check.py
│ │ ├── dependency_versions_table.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ └── rl/
│ │ │ ├── __init__.py
│ │ │ └── value_guided_sampling.py
│ │ ├── image_processor.py
│ │ ├── loaders.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ ├── activations.py
│ │ │ ├── attention.py
│ │ │ ├── attention_flax.py
│ │ │ ├── attention_processor.py
│ │ │ ├── autoencoder_kl.py
│ │ │ ├── controlnet.py
│ │ │ ├── controlnet_flax.py
│ │ │ ├── cross_attention.py
│ │ │ ├── dual_transformer_2d.py
│ │ │ ├── embeddings.py
│ │ │ ├── embeddings_flax.py
│ │ │ ├── modeling_flax_pytorch_utils.py
│ │ │ ├── modeling_flax_utils.py
│ │ │ ├── modeling_pytorch_flax_utils.py
│ │ │ ├── modeling_utils.py
│ │ │ ├── prior_transformer.py
│ │ │ ├── resnet.py
│ │ │ ├── resnet_flax.py
│ │ │ ├── t5_film_transformer.py
│ │ │ ├── transformer_2d.py
│ │ │ ├── transformer_temporal.py
│ │ │ ├── unet_1d.py
│ │ │ ├── unet_1d_blocks.py
│ │ │ ├── unet_2d.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ ├── unet_2d_blocks_flax.py
│ │ │ ├── unet_2d_condition.py
│ │ │ ├── unet_2d_condition_flax.py
│ │ │ ├── unet_3d_blocks.py
│ │ │ ├── unet_3d_condition.py
│ │ │ ├── vae.py
│ │ │ ├── vae_flax.py
│ │ │ └── vq_model.py
│ │ ├── optimization.py
│ │ ├── pipeline_utils.py
│ │ ├── pipelines/
│ │ │ ├── __init__.py
│ │ │ ├── alt_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_roberta_series.py
│ │ │ │ ├── pipeline_alt_diffusion.py
│ │ │ │ └── pipeline_alt_diffusion_img2img.py
│ │ │ ├── audio_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mel.py
│ │ │ │ └── pipeline_audio_diffusion.py
│ │ │ ├── audioldm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_audioldm.py
│ │ │ ├── consistency_models/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_consistency_models.py
│ │ │ ├── controlnet/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── multicontrolnet.py
│ │ │ │ ├── pipeline_controlnet.py
│ │ │ │ ├── pipeline_controlnet_img2img.py
│ │ │ │ ├── pipeline_controlnet_inpaint.py
│ │ │ │ └── pipeline_flax_controlnet.py
│ │ │ ├── dance_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dance_diffusion.py
│ │ │ ├── ddim/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddim.py
│ │ │ ├── ddpm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddpm.py
│ │ │ ├── deepfloyd_if/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_if.py
│ │ │ │ ├── pipeline_if_img2img.py
│ │ │ │ ├── pipeline_if_img2img_superresolution.py
│ │ │ │ ├── pipeline_if_inpainting.py
│ │ │ │ ├── pipeline_if_inpainting_superresolution.py
│ │ │ │ ├── pipeline_if_superresolution.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── timesteps.py
│ │ │ │ └── watermark.py
│ │ │ ├── dit/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dit.py
│ │ │ ├── kandinsky/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky.py
│ │ │ │ ├── pipeline_kandinsky_img2img.py
│ │ │ │ ├── pipeline_kandinsky_inpaint.py
│ │ │ │ ├── pipeline_kandinsky_prior.py
│ │ │ │ └── text_encoder.py
│ │ │ ├── kandinsky2_2/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky2_2.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_inpainting.py
│ │ │ │ ├── pipeline_kandinsky2_2_prior.py
│ │ │ │ └── pipeline_kandinsky2_2_prior_emb2emb.py
│ │ │ ├── latent_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_latent_diffusion.py
│ │ │ │ └── pipeline_latent_diffusion_superresolution.py
│ │ │ ├── latent_diffusion_uncond/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_latent_diffusion_uncond.py
│ │ │ ├── onnx_utils.py
│ │ │ ├── paint_by_example/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── pipeline_paint_by_example.py
│ │ │ ├── pipeline_flax_utils.py
│ │ │ ├── pipeline_utils.py
│ │ │ ├── pndm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_pndm.py
│ │ │ ├── repaint/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_repaint.py
│ │ │ ├── score_sde_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_score_sde_ve.py
│ │ │ ├── semantic_stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_semantic_stable_diffusion.py
│ │ │ ├── shap_e/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── camera.py
│ │ │ │ ├── pipeline_shap_e.py
│ │ │ │ ├── pipeline_shap_e_img2img.py
│ │ │ │ └── renderer.py
│ │ │ ├── spectrogram_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── continous_encoder.py
│ │ │ │ ├── midi_utils.py
│ │ │ │ ├── notes_encoder.py
│ │ │ │ └── pipeline_spectrogram_diffusion.py
│ │ │ ├── stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── convert_from_ckpt.py
│ │ │ │ ├── pipeline_cycle_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_attend_and_excite.py
│ │ │ │ ├── pipeline_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_stable_diffusion_depth2img.py
│ │ │ │ ├── pipeline_stable_diffusion_diffedit.py
│ │ │ │ ├── pipeline_stable_diffusion_image_variation.py
│ │ │ │ ├── pipeline_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_stable_diffusion_instruct_pix2pix.py
│ │ │ │ ├── pipeline_stable_diffusion_k_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_latent_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion_ldm3d.py
│ │ │ │ ├── pipeline_stable_diffusion_model_editing.py
│ │ │ │ ├── pipeline_stable_diffusion_panorama.py
│ │ │ │ ├── pipeline_stable_diffusion_paradigms.py
│ │ │ │ ├── pipeline_stable_diffusion_pix2pix_zero.py
│ │ │ │ ├── pipeline_stable_diffusion_sag.py
│ │ │ │ ├── pipeline_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_unclip.py
│ │ │ │ ├── pipeline_stable_unclip_img2img.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── safety_checker_flax.py
│ │ │ │ └── stable_unclip_image_normalizer.py
│ │ │ ├── stable_diffusion_safe/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_safe.py
│ │ │ │ └── safety_checker.py
│ │ │ ├── stable_diffusion_xl/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_xl.py
│ │ │ │ ├── pipeline_stable_diffusion_xl_img2img.py
│ │ │ │ └── watermark.py
│ │ │ ├── stochastic_karras_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_stochastic_karras_ve.py
│ │ │ ├── text_to_video_synthesis/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_text_to_video_synth.py
│ │ │ │ ├── pipeline_text_to_video_synth_img2img.py
│ │ │ │ └── pipeline_text_to_video_zero.py
│ │ │ ├── unclip/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_unclip.py
│ │ │ │ ├── pipeline_unclip_image_variation.py
│ │ │ │ └── text_proj.py
│ │ │ ├── unidiffuser/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_decoder.py
│ │ │ │ ├── modeling_uvit.py
│ │ │ │ └── pipeline_unidiffuser.py
│ │ │ ├── versatile_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_unet.py
│ │ │ │ ├── pipeline_versatile_diffusion.py
│ │ │ │ ├── pipeline_versatile_diffusion_dual_guided.py
│ │ │ │ ├── pipeline_versatile_diffusion_image_variation.py
│ │ │ │ └── pipeline_versatile_diffusion_text_to_image.py
│ │ │ └── vq_diffusion/
│ │ │ ├── __init__.py
│ │ │ └── pipeline_vq_diffusion.py
│ │ ├── schedulers/
│ │ │ ├── __init__.py
│ │ │ ├── scheduling_consistency_models.py
│ │ │ ├── scheduling_ddim.py
│ │ │ ├── scheduling_ddim_flax.py
│ │ │ ├── scheduling_ddim_inverse.py
│ │ │ ├── scheduling_ddim_parallel.py
│ │ │ ├── scheduling_ddpm.py
│ │ │ ├── scheduling_ddpm_flax.py
│ │ │ ├── scheduling_ddpm_parallel.py
│ │ │ ├── scheduling_deis_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep_flax.py
│ │ │ ├── scheduling_dpmsolver_multistep_inverse.py
│ │ │ ├── scheduling_dpmsolver_sde.py
│ │ │ ├── scheduling_dpmsolver_singlestep.py
│ │ │ ├── scheduling_euler_ancestral_discrete.py
│ │ │ ├── scheduling_euler_discrete.py
│ │ │ ├── scheduling_heun_discrete.py
│ │ │ ├── scheduling_ipndm.py
│ │ │ ├── scheduling_k_dpm_2_ancestral_discrete.py
│ │ │ ├── scheduling_k_dpm_2_discrete.py
│ │ │ ├── scheduling_karras_ve.py
│ │ │ ├── scheduling_karras_ve_flax.py
│ │ │ ├── scheduling_lms_discrete.py
│ │ │ ├── scheduling_lms_discrete_flax.py
│ │ │ ├── scheduling_pndm.py
│ │ │ ├── scheduling_pndm_flax.py
│ │ │ ├── scheduling_repaint.py
│ │ │ ├── scheduling_sde_ve.py
│ │ │ ├── scheduling_sde_ve_flax.py
│ │ │ ├── scheduling_sde_vp.py
│ │ │ ├── scheduling_unclip.py
│ │ │ ├── scheduling_unipc_multistep.py
│ │ │ ├── scheduling_utils.py
│ │ │ ├── scheduling_utils_flax.py
│ │ │ └── scheduling_vq_diffusion.py
│ │ ├── training_utils.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── accelerate_utils.py
│ │ ├── constants.py
│ │ ├── deprecation_utils.py
│ │ ├── doc_utils.py
│ │ ├── dummy_flax_and_transformers_objects.py
│ │ ├── dummy_flax_objects.py
│ │ ├── dummy_note_seq_objects.py
│ │ ├── dummy_onnx_objects.py
│ │ ├── dummy_pt_objects.py
│ │ ├── dummy_torch_and_librosa_objects.py
│ │ ├── dummy_torch_and_scipy_objects.py
│ │ ├── dummy_torch_and_torchsde_objects.py
│ │ ├── dummy_torch_and_transformers_and_invisible_watermark_objects.py
│ │ ├── dummy_torch_and_transformers_and_k_diffusion_objects.py
│ │ ├── dummy_torch_and_transformers_and_onnx_objects.py
│ │ ├── dummy_torch_and_transformers_objects.py
│ │ ├── dummy_transformers_and_torch_and_note_seq_objects.py
│ │ ├── dynamic_modules_utils.py
│ │ ├── hub_utils.py
│ │ ├── import_utils.py
│ │ ├── logging.py
│ │ ├── model_card_template.md
│ │ ├── outputs.py
│ │ ├── pil_utils.py
│ │ ├── testing_utils.py
│ │ └── torch_utils.py
│ ├── pipeline_zero1to3.py
│ ├── train_eschernet.py
│ └── unet_2d_condition.py
├── 6DoF/
│ ├── CN_encoder.py
│ ├── dataset.py
│ ├── diffusers/
│ │ ├── __init__.py
│ │ ├── commands/
│ │ │ ├── __init__.py
│ │ │ ├── diffusers_cli.py
│ │ │ └── env.py
│ │ ├── configuration_utils.py
│ │ ├── dependency_versions_check.py
│ │ ├── dependency_versions_table.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ └── rl/
│ │ │ ├── __init__.py
│ │ │ └── value_guided_sampling.py
│ │ ├── image_processor.py
│ │ ├── loaders.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ ├── activations.py
│ │ │ ├── attention.py
│ │ │ ├── attention_flax.py
│ │ │ ├── attention_processor.py
│ │ │ ├── autoencoder_kl.py
│ │ │ ├── controlnet.py
│ │ │ ├── controlnet_flax.py
│ │ │ ├── cross_attention.py
│ │ │ ├── dual_transformer_2d.py
│ │ │ ├── embeddings.py
│ │ │ ├── embeddings_flax.py
│ │ │ ├── modeling_flax_pytorch_utils.py
│ │ │ ├── modeling_flax_utils.py
│ │ │ ├── modeling_pytorch_flax_utils.py
│ │ │ ├── modeling_utils.py
│ │ │ ├── prior_transformer.py
│ │ │ ├── resnet.py
│ │ │ ├── resnet_flax.py
│ │ │ ├── t5_film_transformer.py
│ │ │ ├── transformer_2d.py
│ │ │ ├── transformer_temporal.py
│ │ │ ├── unet_1d.py
│ │ │ ├── unet_1d_blocks.py
│ │ │ ├── unet_2d.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ ├── unet_2d_blocks_flax.py
│ │ │ ├── unet_2d_condition.py
│ │ │ ├── unet_2d_condition_flax.py
│ │ │ ├── unet_3d_blocks.py
│ │ │ ├── unet_3d_condition.py
│ │ │ ├── vae.py
│ │ │ ├── vae_flax.py
│ │ │ └── vq_model.py
│ │ ├── optimization.py
│ │ ├── pipeline_utils.py
│ │ ├── pipelines/
│ │ │ ├── __init__.py
│ │ │ ├── alt_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_roberta_series.py
│ │ │ │ ├── pipeline_alt_diffusion.py
│ │ │ │ └── pipeline_alt_diffusion_img2img.py
│ │ │ ├── audio_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mel.py
│ │ │ │ └── pipeline_audio_diffusion.py
│ │ │ ├── audioldm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_audioldm.py
│ │ │ ├── consistency_models/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_consistency_models.py
│ │ │ ├── controlnet/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── multicontrolnet.py
│ │ │ │ ├── pipeline_controlnet.py
│ │ │ │ ├── pipeline_controlnet_img2img.py
│ │ │ │ ├── pipeline_controlnet_inpaint.py
│ │ │ │ └── pipeline_flax_controlnet.py
│ │ │ ├── dance_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dance_diffusion.py
│ │ │ ├── ddim/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddim.py
│ │ │ ├── ddpm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_ddpm.py
│ │ │ ├── deepfloyd_if/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_if.py
│ │ │ │ ├── pipeline_if_img2img.py
│ │ │ │ ├── pipeline_if_img2img_superresolution.py
│ │ │ │ ├── pipeline_if_inpainting.py
│ │ │ │ ├── pipeline_if_inpainting_superresolution.py
│ │ │ │ ├── pipeline_if_superresolution.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── timesteps.py
│ │ │ │ └── watermark.py
│ │ │ ├── dit/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_dit.py
│ │ │ ├── kandinsky/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky.py
│ │ │ │ ├── pipeline_kandinsky_img2img.py
│ │ │ │ ├── pipeline_kandinsky_inpaint.py
│ │ │ │ ├── pipeline_kandinsky_prior.py
│ │ │ │ └── text_encoder.py
│ │ │ ├── kandinsky2_2/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_kandinsky2_2.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet.py
│ │ │ │ ├── pipeline_kandinsky2_2_controlnet_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_img2img.py
│ │ │ │ ├── pipeline_kandinsky2_2_inpainting.py
│ │ │ │ ├── pipeline_kandinsky2_2_prior.py
│ │ │ │ └── pipeline_kandinsky2_2_prior_emb2emb.py
│ │ │ ├── latent_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_latent_diffusion.py
│ │ │ │ └── pipeline_latent_diffusion_superresolution.py
│ │ │ ├── latent_diffusion_uncond/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_latent_diffusion_uncond.py
│ │ │ ├── onnx_utils.py
│ │ │ ├── paint_by_example/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── pipeline_paint_by_example.py
│ │ │ ├── pipeline_flax_utils.py
│ │ │ ├── pipeline_utils.py
│ │ │ ├── pndm/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_pndm.py
│ │ │ ├── repaint/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_repaint.py
│ │ │ ├── score_sde_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_score_sde_ve.py
│ │ │ ├── semantic_stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_semantic_stable_diffusion.py
│ │ │ ├── shap_e/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── camera.py
│ │ │ │ ├── pipeline_shap_e.py
│ │ │ │ ├── pipeline_shap_e_img2img.py
│ │ │ │ └── renderer.py
│ │ │ ├── spectrogram_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── continous_encoder.py
│ │ │ │ ├── midi_utils.py
│ │ │ │ ├── notes_encoder.py
│ │ │ │ └── pipeline_spectrogram_diffusion.py
│ │ │ ├── stable_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── convert_from_ckpt.py
│ │ │ │ ├── pipeline_cycle_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_flax_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_onnx_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_attend_and_excite.py
│ │ │ │ ├── pipeline_stable_diffusion_controlnet.py
│ │ │ │ ├── pipeline_stable_diffusion_depth2img.py
│ │ │ │ ├── pipeline_stable_diffusion_diffedit.py
│ │ │ │ ├── pipeline_stable_diffusion_image_variation.py
│ │ │ │ ├── pipeline_stable_diffusion_img2img.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint.py
│ │ │ │ ├── pipeline_stable_diffusion_inpaint_legacy.py
│ │ │ │ ├── pipeline_stable_diffusion_instruct_pix2pix.py
│ │ │ │ ├── pipeline_stable_diffusion_k_diffusion.py
│ │ │ │ ├── pipeline_stable_diffusion_latent_upscale.py
│ │ │ │ ├── pipeline_stable_diffusion_ldm3d.py
│ │ │ │ ├── pipeline_stable_diffusion_model_editing.py
│ │ │ │ ├── pipeline_stable_diffusion_panorama.py
│ │ │ │ ├── pipeline_stable_diffusion_paradigms.py
│ │ │ │ ├── pipeline_stable_diffusion_pix2pix_zero.py
│ │ │ │ ├── pipeline_stable_diffusion_sag.py
│ │ │ │ ├── pipeline_stable_diffusion_upscale.py
│ │ │ │ ├── pipeline_stable_unclip.py
│ │ │ │ ├── pipeline_stable_unclip_img2img.py
│ │ │ │ ├── safety_checker.py
│ │ │ │ ├── safety_checker_flax.py
│ │ │ │ └── stable_unclip_image_normalizer.py
│ │ │ ├── stable_diffusion_safe/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_safe.py
│ │ │ │ └── safety_checker.py
│ │ │ ├── stable_diffusion_xl/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_stable_diffusion_xl.py
│ │ │ │ ├── pipeline_stable_diffusion_xl_img2img.py
│ │ │ │ └── watermark.py
│ │ │ ├── stochastic_karras_ve/
│ │ │ │ ├── __init__.py
│ │ │ │ └── pipeline_stochastic_karras_ve.py
│ │ │ ├── text_to_video_synthesis/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_text_to_video_synth.py
│ │ │ │ ├── pipeline_text_to_video_synth_img2img.py
│ │ │ │ └── pipeline_text_to_video_zero.py
│ │ │ ├── unclip/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── pipeline_unclip.py
│ │ │ │ ├── pipeline_unclip_image_variation.py
│ │ │ │ └── text_proj.py
│ │ │ ├── unidiffuser/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_decoder.py
│ │ │ │ ├── modeling_uvit.py
│ │ │ │ └── pipeline_unidiffuser.py
│ │ │ ├── versatile_diffusion/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── modeling_text_unet.py
│ │ │ │ ├── pipeline_versatile_diffusion.py
│ │ │ │ ├── pipeline_versatile_diffusion_dual_guided.py
│ │ │ │ ├── pipeline_versatile_diffusion_image_variation.py
│ │ │ │ └── pipeline_versatile_diffusion_text_to_image.py
│ │ │ └── vq_diffusion/
│ │ │ ├── __init__.py
│ │ │ └── pipeline_vq_diffusion.py
│ │ ├── schedulers/
│ │ │ ├── __init__.py
│ │ │ ├── scheduling_consistency_models.py
│ │ │ ├── scheduling_ddim.py
│ │ │ ├── scheduling_ddim_flax.py
│ │ │ ├── scheduling_ddim_inverse.py
│ │ │ ├── scheduling_ddim_parallel.py
│ │ │ ├── scheduling_ddpm.py
│ │ │ ├── scheduling_ddpm_flax.py
│ │ │ ├── scheduling_ddpm_parallel.py
│ │ │ ├── scheduling_deis_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep.py
│ │ │ ├── scheduling_dpmsolver_multistep_flax.py
│ │ │ ├── scheduling_dpmsolver_multistep_inverse.py
│ │ │ ├── scheduling_dpmsolver_sde.py
│ │ │ ├── scheduling_dpmsolver_singlestep.py
│ │ │ ├── scheduling_euler_ancestral_discrete.py
│ │ │ ├── scheduling_euler_discrete.py
│ │ │ ├── scheduling_heun_discrete.py
│ │ │ ├── scheduling_ipndm.py
│ │ │ ├── scheduling_k_dpm_2_ancestral_discrete.py
│ │ │ ├── scheduling_k_dpm_2_discrete.py
│ │ │ ├── scheduling_karras_ve.py
│ │ │ ├── scheduling_karras_ve_flax.py
│ │ │ ├── scheduling_lms_discrete.py
│ │ │ ├── scheduling_lms_discrete_flax.py
│ │ │ ├── scheduling_pndm.py
│ │ │ ├── scheduling_pndm_flax.py
│ │ │ ├── scheduling_repaint.py
│ │ │ ├── scheduling_sde_ve.py
│ │ │ ├── scheduling_sde_ve_flax.py
│ │ │ ├── scheduling_sde_vp.py
│ │ │ ├── scheduling_unclip.py
│ │ │ ├── scheduling_unipc_multistep.py
│ │ │ ├── scheduling_utils.py
│ │ │ ├── scheduling_utils_flax.py
│ │ │ └── scheduling_vq_diffusion.py
│ │ ├── training_utils.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── accelerate_utils.py
│ │ ├── constants.py
│ │ ├── deprecation_utils.py
│ │ ├── doc_utils.py
│ │ ├── dummy_flax_and_transformers_objects.py
│ │ ├── dummy_flax_objects.py
│ │ ├── dummy_note_seq_objects.py
│ │ ├── dummy_onnx_objects.py
│ │ ├── dummy_pt_objects.py
│ │ ├── dummy_torch_and_librosa_objects.py
│ │ ├── dummy_torch_and_scipy_objects.py
│ │ ├── dummy_torch_and_torchsde_objects.py
│ │ ├── dummy_torch_and_transformers_and_invisible_watermark_objects.py
│ │ ├── dummy_torch_and_transformers_and_k_diffusion_objects.py
│ │ ├── dummy_torch_and_transformers_and_onnx_objects.py
│ │ ├── dummy_torch_and_transformers_objects.py
│ │ ├── dummy_transformers_and_torch_and_note_seq_objects.py
│ │ ├── dynamic_modules_utils.py
│ │ ├── hub_utils.py
│ │ ├── import_utils.py
│ │ ├── logging.py
│ │ ├── model_card_template.md
│ │ ├── outputs.py
│ │ ├── pil_utils.py
│ │ ├── testing_utils.py
│ │ └── torch_utils.py
│ ├── instructions.md
│ ├── pipeline_zero1to3.py
│ ├── train_eschernet.py
│ └── unet_2d_condition.py
├── CaPE.py
├── LICENSE
├── README.md
├── demo/
│ └── GSO30/
│ └── elephant/
│ └── render_mvs_25/
│ └── model/
│ ├── 000.npy
│ ├── 001.npy
│ ├── 002.npy
│ ├── 003.npy
│ ├── 004.npy
│ ├── 005.npy
│ ├── 006.npy
│ ├── 007.npy
│ ├── 008.npy
│ ├── 009.npy
│ ├── 010.npy
│ ├── 011.npy
│ ├── 012.npy
│ ├── 013.npy
│ ├── 014.npy
│ ├── 015.npy
│ ├── 016.npy
│ ├── 017.npy
│ ├── 018.npy
│ ├── 019.npy
│ ├── 020.npy
│ ├── 021.npy
│ ├── 022.npy
│ ├── 023.npy
│ └── 024.npy
├── environment.yml
├── eval_eschernet.py
├── eval_eschernet.sh
├── metrics/
│ ├── NeRF_idx/
│ │ ├── chair/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── drums/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── ficus/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── hotdog/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── lego/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── materials/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ ├── mic/
│ │ │ ├── test_M20.npy
│ │ │ ├── train_N100M20_random.npy
│ │ │ ├── train_N10M20_random.npy
│ │ │ ├── train_N1M20_random.npy
│ │ │ ├── train_N20M20_random.npy
│ │ │ ├── train_N2M20_random.npy
│ │ │ ├── train_N30M20_random.npy
│ │ │ ├── train_N3M20_random.npy
│ │ │ ├── train_N50M20_random.npy
│ │ │ ├── train_N5M20_random.npy
│ │ │ ├── train_fov0M20_fov.npy
│ │ │ ├── train_fov135M20_fov.npy
│ │ │ ├── train_fov15M20_fov.npy
│ │ │ ├── train_fov180M20_fov.npy
│ │ │ ├── train_fov225M20_fov.npy
│ │ │ ├── train_fov270M20_fov.npy
│ │ │ ├── train_fov30M20_fov.npy
│ │ │ ├── train_fov360M20_fov.npy
│ │ │ ├── train_fov50M20_fov.npy
│ │ │ ├── train_fov5M20_fov.npy
│ │ │ └── train_fov90M20_fov.npy
│ │ └── ship/
│ │ ├── test_M20.npy
│ │ ├── train_N100M20_random.npy
│ │ ├── train_N10M20_random.npy
│ │ ├── train_N1M20_random.npy
│ │ ├── train_N20M20_random.npy
│ │ ├── train_N2M20_random.npy
│ │ ├── train_N30M20_random.npy
│ │ ├── train_N3M20_random.npy
│ │ ├── train_N50M20_random.npy
│ │ ├── train_N5M20_random.npy
│ │ ├── train_fov0M20_fov.npy
│ │ ├── train_fov135M20_fov.npy
│ │ ├── train_fov15M20_fov.npy
│ │ ├── train_fov180M20_fov.npy
│ │ ├── train_fov225M20_fov.npy
│ │ ├── train_fov270M20_fov.npy
│ │ ├── train_fov30M20_fov.npy
│ │ ├── train_fov360M20_fov.npy
│ │ ├── train_fov50M20_fov.npy
│ │ ├── train_fov5M20_fov.npy
│ │ └── train_fov90M20_fov.npy
│ ├── RenderOption_phong.json
│ ├── RenderOption_rgb.json
│ ├── ScreenCamera_0.json
│ ├── ScreenCamera_1.json
│ ├── ScreenCamera_2.json
│ ├── ScreenCamera_3.json
│ ├── ScreenCamera_4.json
│ ├── ScreenCamera_5.json
│ ├── eval_2D_NVS.py
│ ├── eval_3D_GSO.py
│ └── metrics.py
└── scripts/
├── README.md
├── all_invalid.npy
├── blender_script_mvs.py
├── empty_ids.npy
├── invalid_ids.npy
├── obj_ids.npy
├── objaverse_filter.py
└── render_all_mvs.py
Showing preview only (571K chars total). Download the full file or copy to clipboard to get everything.
SYMBOL INDEX (6818 symbols across 470 files)
FILE: 3drecon/raymarching/backend.py
function find_cl_path (line 17) | def find_cl_path():
FILE: 3drecon/raymarching/raymarching.py
class _near_far_from_aabb (line 19) | class _near_far_from_aabb(Function):
method forward (line 22) | def forward(ctx, rays_o, rays_d, aabb, min_near=0.2):
class _sph_from_ray (line 52) | class _sph_from_ray(Function):
method forward (line 55) | def forward(ctx, rays_o, rays_d, radius):
class _morton3D (line 83) | class _morton3D(Function):
method forward (line 85) | def forward(ctx, coords):
class _morton3D_invert (line 106) | class _morton3D_invert(Function):
method forward (line 108) | def forward(ctx, indices):
class _packbits (line 129) | class _packbits(Function):
method forward (line 132) | def forward(ctx, grid, thresh, bitfield=None):
class _march_rays_train (line 161) | class _march_rays_train(Function):
method forward (line 164) | def forward(ctx, rays_o, rays_d, bound, density_bitfield, C, H, nears,...
class _composite_rays_train (line 238) | class _composite_rays_train(Function):
method forward (line 241) | def forward(ctx, sigmas, rgbs, deltas, rays, T_thresh=1e-4):
method backward (line 273) | def backward(ctx, grad_weights_sum, grad_depth, grad_image):
class _march_rays (line 297) | class _march_rays(Function):
method forward (line 300) | def forward(ctx, n_alive, n_step, rays_alive, rays_t, rays_o, rays_d, ...
class _composite_rays (line 351) | class _composite_rays(Function):
method forward (line 354) | def forward(ctx, n_alive, n_step, rays_alive, rays_t, sigmas, rgbs, de...
FILE: 3drecon/raymarching/setup.py
function find_cl_path (line 18) | def find_cl_path():
FILE: 3drecon/raymarching/src/bindings.cpp
function PYBIND11_MODULE (line 5) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: 3drecon/renderer/agg_net.py
function weights_init (line 5) | def weights_init(m):
class NeRF (line 11) | class NeRF(nn.Module):
method __init__ (line 12) | def __init__(self, vol_n=8+8, feat_ch=8+16+32+3, hid_n=64):
method forward (line 27) | def forward(self, vox_feat, img_feat_rgb_dir, source_img_mask):
class Agg (line 44) | class Agg(nn.Module):
method __init__ (line 45) | def __init__(self, feat_ch):
method masked_mean_var (line 58) | def masked_mean_var(self, img_feat_rgb, source_img_mask):
method forward (line 66) | def forward(self, img_feat_rgb_dir, source_img_mask):
FILE: 3drecon/renderer/cost_reg_net.py
class ConvBnReLU3D (line 3) | class ConvBnReLU3D(nn.Module):
method __init__ (line 4) | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,...
method forward (line 10) | def forward(self, x):
class CostRegNet (line 13) | class CostRegNet(nn.Module):
method __init__ (line 14) | def __init__(self, in_channels, norm_act=nn.BatchNorm3d):
method forward (line 44) | def forward(self, x):
class MinCostRegNet (line 60) | class MinCostRegNet(nn.Module):
method __init__ (line 61) | def __init__(self, in_channels, norm_act=nn.BatchNorm3d):
method forward (line 84) | def forward(self, x):
FILE: 3drecon/renderer/dummy_dataset.py
class DummyDataset (line 5) | class DummyDataset(pl.LightningDataModule):
method __init__ (line 6) | def __init__(self,seed):
method setup (line 9) | def setup(self, stage):
method train_dataloader (line 16) | def train_dataloader(self):
method val_dataloader (line 19) | def val_dataloader(self):
method test_dataloader (line 22) | def test_dataloader(self):
class DummyData (line 25) | class DummyData(Dataset):
method __init__ (line 26) | def __init__(self,is_train):
method __len__ (line 29) | def __len__(self):
method __getitem__ (line 35) | def __getitem__(self, index):
FILE: 3drecon/renderer/feature_net.py
class ConvBnReLU (line 4) | class ConvBnReLU(nn.Module):
method __init__ (line 5) | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,...
method forward (line 11) | def forward(self, x):
class FeatureNet (line 14) | class FeatureNet(nn.Module):
method __init__ (line 15) | def __init__(self, norm_act=nn.BatchNorm2d):
method _upsample_add (line 28) | def _upsample_add(self, x, y):
method forward (line 31) | def forward(self, x):
FILE: 3drecon/renderer/neus_networks.py
class Embedder (line 10) | class Embedder:
method __init__ (line 11) | def __init__(self, **kwargs):
method create_embedding_fn (line 15) | def create_embedding_fn(self):
method embed (line 39) | def embed(self, inputs):
function get_embedder (line 43) | def get_embedder(multires, input_dims=3):
class SDFNetwork (line 60) | class SDFNetwork(nn.Module):
method __init__ (line 61) | def __init__(self, d_in, d_out, d_hidden, n_layers, skip_in=(4,), mult...
method forward (line 113) | def forward(self, inputs):
method sdf (line 132) | def sdf(self, x):
method sdf_hidden_appearance (line 135) | def sdf_hidden_appearance(self, x):
method gradient (line 138) | def gradient(self, x):
method sdf_normal (line 152) | def sdf_normal(self, x):
class SDFNetworkWithFeature (line 166) | class SDFNetworkWithFeature(nn.Module):
method __init__ (line 167) | def __init__(self, cube, dp_in, df_in, d_out, d_hidden, n_layers, skip...
method forward (line 221) | def forward(self, points):
method sdf (line 247) | def sdf(self, x):
method sdf_hidden_appearance (line 250) | def sdf_hidden_appearance(self, x):
method gradient (line 253) | def gradient(self, x):
method sdf_normal (line 267) | def sdf_normal(self, x):
class VanillaMLP (line 282) | class VanillaMLP(nn.Module):
method __init__ (line 283) | def __init__(self, dim_in, dim_out, n_neurons, n_hidden_layers):
method forward (line 295) | def forward(self, x):
method make_linear (line 299) | def make_linear(self, dim_in, dim_out, is_first, is_last):
method make_activation (line 320) | def make_activation(self):
class SDFHashGridNetwork (line 327) | class SDFHashGridNetwork(nn.Module):
method __init__ (line 328) | def __init__(self, bound=0.5, feats_dim=13):
method forward (line 358) | def forward(self, x):
method sdf (line 369) | def sdf(self, x):
method gradient (line 372) | def gradient(self, x):
method sdf_normal (line 386) | def sdf_normal(self, x):
class RenderingFFNetwork (line 400) | class RenderingFFNetwork(nn.Module):
method __init__ (line 401) | def __init__(self, in_feats_dim=12):
method forward (line 422) | def forward(self, points, normals, view_dirs, feature_vectors):
class RenderingNetwork (line 433) | class RenderingNetwork(nn.Module):
method __init__ (line 434) | def __init__(self, d_feature, d_in, d_out, d_hidden,
method forward (line 463) | def forward(self, points, normals, view_dirs, feature_vectors):
class SingleVarianceNetwork (line 488) | class SingleVarianceNetwork(nn.Module):
method __init__ (line 489) | def __init__(self, init_val, activation='exp'):
method forward (line 494) | def forward(self, x):
method warp (line 501) | def warp(self, x, inv_s):
FILE: 3drecon/renderer/ngp_renderer.py
function custom_meshgrid (line 17) | def custom_meshgrid(*args):
function sample_pdf (line 24) | def sample_pdf(bins, weights, n_samples, det=False):
function plot_pointcloud (line 61) | def plot_pointcloud(pc, color=None):
class NGPRenderer (line 73) | class NGPRenderer(nn.Module):
method __init__ (line 74) | def __init__(self,
method forward (line 115) | def forward(self, x, d):
method density (line 119) | def density(self, x):
method color (line 122) | def color(self, x, d, mask=None, **kwargs):
method reset_extra_state (line 125) | def reset_extra_state(self):
method run (line 137) | def run(self, rays_o, rays_d, num_steps=128, upsample_steps=128, bg_co...
method run_cuda (line 268) | def run_cuda(self, rays_o, rays_d, dt_gamma=0, bg_color=None, perturb=...
method mark_untrained_grid (line 369) | def mark_untrained_grid(self, poses, intrinsic, S=64):
method update_extra_state (line 434) | def update_extra_state(self, decay=0.95, S=128):
method render (line 529) | def render(self, rays_o, rays_d, staged=False, max_ray_batch=4096, **k...
class _trunc_exp (line 543) | class _trunc_exp(Function):
method forward (line 546) | def forward(ctx, x):
method backward (line 552) | def backward(ctx, g):
class NGPNetwork (line 558) | class NGPNetwork(NGPRenderer):
method __init__ (line 559) | def __init__(self,
method forward (line 636) | def forward(self, x, d):
method density (line 669) | def density(self, x):
method color (line 691) | def color(self, x, d, mask=None, geo_feat=None, **kwargs):
FILE: 3drecon/renderer/renderer.py
function sample_pdf (line 24) | def sample_pdf(bins, weights, n_samples, det=True):
function near_far_from_sphere (line 58) | def near_far_from_sphere(rays_o, rays_d, radius=DEFAULT_RADIUS):
class BackgroundRemoval (line 66) | class BackgroundRemoval:
method __init__ (line 67) | def __init__(self, device='cuda'):
method __call__ (line 83) | def __call__(self, image):
class BaseRenderer (line 91) | class BaseRenderer(nn.Module):
method __init__ (line 92) | def __init__(self, train_batch_num, test_batch_num):
method render_impl (line 98) | def render_impl(self, ray_batch, is_train, step):
method render_with_loss (line 102) | def render_with_loss(self, ray_batch, is_train, step):
method render (line 105) | def render(self, ray_batch, is_train, step):
class NeuSRenderer (line 123) | class NeuSRenderer(BaseRenderer):
method __init__ (line 124) | def __init__(self, train_batch_num, test_batch_num, lambda_eikonal_los...
method get_vertex_colors (line 143) | def get_vertex_colors(self, vertices):
method upsample (line 164) | def upsample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
method cat_z_vals (line 197) | def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=Fal...
method sample_depth (line 214) | def sample_depth(self, rays_o, rays_d, near, far, perturb):
method compute_sdf_alpha (line 242) | def compute_sdf_alpha(self, points, dists, dirs, cos_anneal_ratio, step):
method get_anneal_val (line 269) | def get_anneal_val(self, step):
method get_inner_mask (line 275) | def get_inner_mask(self, points):
method render_impl (line 278) | def render_impl(self, ray_batch, is_train, step):
method render_with_loss (line 322) | def render_with_loss(self, ray_batch, is_train, step):
class NeRFRenderer (line 350) | class NeRFRenderer(BaseRenderer):
method __init__ (line 351) | def __init__(self, train_batch_num, test_batch_num, bound=0.5, use_mas...
method render_impl (line 363) | def render_impl(self, ray_batch, is_train, step):
method render_with_loss (line 378) | def render_with_loss(self, ray_batch, is_train, step):
function cartesian_to_spherical (line 396) | def cartesian_to_spherical(xyz):
function get_pose (line 405) | def get_pose(target_RT):
class RendererTrainer (line 412) | class RendererTrainer(pl.LightningModule):
method __init__ (line 413) | def __init__(self, image_path, data_path, total_steps, warm_up_steps, ...
method _construct_ray_batch (line 453) | def _construct_ray_batch(self, images_info):
method load_model (line 479) | def load_model(cfg, ckpt):
method _init_dataset (line 488) | def _init_dataset(self):
method _shuffle_train_batch (line 548) | def _shuffle_train_batch(self):
method _shuffle_train_fg_batch (line 554) | def _shuffle_train_fg_batch(self):
method training_step (line 561) | def training_step(self, batch, batch_idx):
method _slice_images_info (line 581) | def _slice_images_info(self, index):
method validation_step (line 585) | def validation_step(self, batch, batch_idx):
method configure_optimizers (line 612) | def configure_optimizers(self):
FILE: 3drecon/train_renderer.py
class ResumeCallBacks (line 22) | class ResumeCallBacks(Callback):
method __init__ (line 23) | def __init__(self):
method on_train_start (line 26) | def on_train_start(self, trainer, pl_module):
function render_images (line 29) | def render_images(model, output,):
function extract_fields (line 77) | def extract_fields(bound_min, bound_max, resolution, query_func, batch_s...
function extract_geometry (line 97) | def extract_geometry(bound_min, bound_max, resolution, threshold, query_...
function extract_mesh (line 107) | def extract_mesh(model, output, resolution=512):
function main (line 118) | def main():
FILE: 3drecon/util.py
function instantiate_from_config (line 6) | def instantiate_from_config(config):
function get_obj_from_str (line 16) | def get_obj_from_str(string, reload=False):
function read_pickle (line 23) | def read_pickle(pkl_path):
function output_points (line 27) | def output_points(fn,pts,colors=None):
function concat_images (line 35) | def concat_images(img0,img1,vert=False):
function concat_images_list (line 49) | def concat_images_list(*args,vert=False):
FILE: 4DoF/CN_encoder.py
class CN_encoder (line 6) | class CN_encoder(ConvNextV2Model):
method __init__ (line 7) | def __init__(self, config):
method forward (line 10) | def forward(
FILE: 4DoF/dataset.py
class ObjaverseDataLoader (line 15) | class ObjaverseDataLoader():
method __init__ (line 16) | def __init__(self, root_dir, batch_size, total_view=12, num_workers=4):
method train_dataloader (line 27) | def train_dataloader(self):
method val_dataloader (line 34) | def val_dataloader(self):
function cartesian_to_spherical (line 40) | def cartesian_to_spherical(xyz):
function get_pose (line 49) | def get_pose(target_RT):
class ObjaverseData (line 63) | class ObjaverseData(Dataset):
method __init__ (line 64) | def __init__(self,
method __len__ (line 106) | def __len__(self):
method cartesian_to_spherical (line 109) | def cartesian_to_spherical(self, xyz):
method get_T (line 118) | def get_T(self, target_RT, cond_RT):
method get_pose (line 135) | def get_pose(self, target_RT):
method load_im (line 148) | def load_im(self, path, color):
method __getitem__ (line 161) | def __getitem__(self, index):
method process_im (line 226) | def process_im(self, im):
FILE: 4DoF/diffusers/commands/__init__.py
class BaseDiffusersCLICommand (line 19) | class BaseDiffusersCLICommand(ABC):
method register_subcommand (line 22) | def register_subcommand(parser: ArgumentParser):
method run (line 26) | def run(self):
FILE: 4DoF/diffusers/commands/diffusers_cli.py
function main (line 21) | def main():
FILE: 4DoF/diffusers/commands/env.py
function info_command_factory (line 25) | def info_command_factory(_):
class EnvironmentCommand (line 29) | class EnvironmentCommand(BaseDiffusersCLICommand):
method register_subcommand (line 31) | def register_subcommand(parser: ArgumentParser):
method run (line 35) | def run(self):
method format_dict (line 83) | def format_dict(d):
FILE: 4DoF/diffusers/configuration_utils.py
class FrozenDict (line 50) | class FrozenDict(OrderedDict):
method __init__ (line 51) | def __init__(self, *args, **kwargs):
method __delitem__ (line 59) | def __delitem__(self, *args, **kwargs):
method setdefault (line 62) | def setdefault(self, *args, **kwargs):
method pop (line 65) | def pop(self, *args, **kwargs):
method update (line 68) | def update(self, *args, **kwargs):
method __setattr__ (line 71) | def __setattr__(self, name, value):
method __setitem__ (line 76) | def __setitem__(self, name, value):
class ConfigMixin (line 82) | class ConfigMixin:
method register_to_config (line 104) | def register_to_config(self, **kwargs):
method __getattr__ (line 121) | def __getattr__(self, name: str) -> Any:
method save_config (line 139) | def save_config(self, save_directory: Union[str, os.PathLike], push_to...
method from_config (line 160) | def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None,...
method get_config_dict (line 246) | def get_config_dict(cls, *args, **kwargs):
method load_config (line 255) | def load_config(
method _get_init_keys (line 421) | def _get_init_keys(cls):
method extract_init_dict (line 425) | def extract_init_dict(cls, config_dict, **kwargs):
method _dict_from_json_file (line 513) | def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):
method __repr__ (line 518) | def __repr__(self):
method config (line 522) | def config(self) -> Dict[str, Any]:
method to_json_string (line 531) | def to_json_string(self) -> str:
method to_json_file (line 557) | def to_json_file(self, json_file_path: Union[str, os.PathLike]):
function register_to_config (line 569) | def register_to_config(init):
function flax_register_to_config (line 619) | def flax_register_to_config(cls):
FILE: 4DoF/diffusers/dependency_versions_check.py
function dep_version_check (line 46) | def dep_version_check(pkg, hint=None):
FILE: 4DoF/diffusers/experimental/rl/value_guided_sampling.py
class ValueGuidedRLPipeline (line 25) | class ValueGuidedRLPipeline(DiffusionPipeline):
method __init__ (line 42) | def __init__(
method normalize (line 70) | def normalize(self, x_in, key):
method de_normalize (line 73) | def de_normalize(self, x_in, key):
method to_torch (line 76) | def to_torch(self, x_in):
method reset_x0 (line 83) | def reset_x0(self, x_in, cond, act_dim):
method run_diffusion (line 88) | def run_diffusion(self, x, conditions, n_guide_steps, scale):
method __call__ (line 121) | def __call__(self, obs, batch_size=64, planning_horizon=32, n_guide_st...
FILE: 4DoF/diffusers/image_processor.py
class VaeImageProcessor (line 27) | class VaeImageProcessor(ConfigMixin):
method __init__ (line 48) | def __init__(
method numpy_to_pil (line 59) | def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
method pil_to_numpy (line 75) | def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]...
method numpy_to_pt (line 87) | def numpy_to_pt(images: np.ndarray) -> torch.FloatTensor:
method pt_to_numpy (line 98) | def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray:
method normalize (line 106) | def normalize(images):
method denormalize (line 113) | def denormalize(images):
method convert_to_rgb (line 120) | def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:
method resize (line 127) | def resize(
method preprocess (line 147) | def preprocess(
method postprocess (line 215) | def postprocess(
class VaeImageProcessorLDM3D (line 255) | class VaeImageProcessorLDM3D(VaeImageProcessor):
method __init__ (line 273) | def __init__(
method numpy_to_pil (line 283) | def numpy_to_pil(images):
method rgblike_to_depthmap (line 299) | def rgblike_to_depthmap(image):
method numpy_to_depth (line 309) | def numpy_to_depth(self, images):
method postprocess (line 329) | def postprocess(
FILE: 4DoF/diffusers/loaders.py
class AttnProcsLayers (line 70) | class AttnProcsLayers(torch.nn.Module):
method __init__ (line 71) | def __init__(self, state_dict: Dict[str, torch.Tensor]):
class UNet2DConditionLoadersMixin (line 112) | class UNet2DConditionLoadersMixin:
method load_attn_procs (line 116) | def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union...
method save_attn_procs (line 329) | def save_attn_procs(
class TextualInversionLoaderMixin (line 407) | class TextualInversionLoaderMixin:
method maybe_convert_prompt (line 412) | def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenize...
method _maybe_convert_prompt (line 439) | def _maybe_convert_prompt(self, prompt: str, tokenizer: "PreTrainedTok...
method load_textual_inversion (line 469) | def load_textual_inversion(
class LoraLoaderMixin (line 738) | class LoraLoaderMixin:
method load_lora_weights (line 746) | def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Uni...
method lora_scale (line 909) | def lora_scale(self) -> float:
method text_encoder_lora_attn_procs (line 915) | def text_encoder_lora_attn_procs(self):
method _remove_text_encoder_monkey_patch (line 920) | def _remove_text_encoder_monkey_patch(self):
method _modify_text_encoder (line 933) | def _modify_text_encoder(self, attn_processors: Dict[str, LoRAAttnProc...
method _lora_attn_processor_attr_to_text_encoder_attr (line 970) | def _lora_attn_processor_attr_to_text_encoder_attr(self):
method _load_text_encoder_attn_procs (line 978) | def _load_text_encoder_attn_procs(
method save_lora_weights (line 1151) | def save_lora_weights(
method _convert_kohya_lora_to_diffusers (line 1228) | def _convert_kohya_lora_to_diffusers(self, state_dict):
class FromSingleFileMixin (line 1279) | class FromSingleFileMixin:
method from_ckpt (line 1285) | def from_ckpt(cls, *args, **kwargs):
method from_single_file (line 1291) | def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
FILE: 4DoF/diffusers/models/activations.py
function get_activation (line 4) | def get_activation(act_fn):
FILE: 4DoF/diffusers/models/attention.py
class BasicTransformerBlock (line 27) | class BasicTransformerBlock(nn.Module):
method __init__ (line 48) | def __init__(
method set_chunk_feed_forward (line 126) | def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
method forward (line 131) | def forward(
class FeedForward (line 210) | class FeedForward(nn.Module):
method __init__ (line 223) | def __init__(
method forward (line 256) | def forward(self, hidden_states):
class GELU (line 262) | class GELU(nn.Module):
method __init__ (line 267) | def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):
method gelu (line 272) | def gelu(self, gate):
method forward (line 278) | def forward(self, hidden_states):
class GEGLU (line 284) | class GEGLU(nn.Module):
method __init__ (line 293) | def __init__(self, dim_in: int, dim_out: int):
method gelu (line 297) | def gelu(self, gate):
method forward (line 303) | def forward(self, hidden_states):
class ApproximateGELU (line 308) | class ApproximateGELU(nn.Module):
method __init__ (line 315) | def __init__(self, dim_in: int, dim_out: int):
method forward (line 319) | def forward(self, x):
class AdaLayerNorm (line 324) | class AdaLayerNorm(nn.Module):
method __init__ (line 329) | def __init__(self, embedding_dim, num_embeddings):
method forward (line 336) | def forward(self, x, timestep):
class AdaLayerNormZero (line 343) | class AdaLayerNormZero(nn.Module):
method __init__ (line 348) | def __init__(self, embedding_dim, num_embeddings):
method forward (line 357) | def forward(self, x, timestep, class_labels, hidden_dtype=None):
class AdaGroupNorm (line 364) | class AdaGroupNorm(nn.Module):
method __init__ (line 369) | def __init__(
method forward (line 383) | def forward(self, x, emb):
FILE: 4DoF/diffusers/models/attention_flax.py
function _query_chunk_attention (line 23) | def _query_chunk_attention(query, key, value, precision, key_chunk_size:...
function jax_memory_efficient_attention (line 74) | def jax_memory_efficient_attention(
class FlaxAttention (line 119) | class FlaxAttention(nn.Module):
method setup (line 145) | def setup(self):
method reshape_heads_to_batch_dim (line 157) | def reshape_heads_to_batch_dim(self, tensor):
method reshape_batch_dim_to_heads (line 165) | def reshape_batch_dim_to_heads(self, tensor):
method __call__ (line 173) | def __call__(self, hidden_states, context=None, deterministic=True):
class FlaxBasicTransformerBlock (line 221) | class FlaxBasicTransformerBlock(nn.Module):
method setup (line 251) | def setup(self):
method __call__ (line 266) | def __call__(self, hidden_states, context, deterministic=True):
class FlaxTransformer2DModel (line 288) | class FlaxTransformer2DModel(nn.Module):
method setup (line 322) | def setup(self):
method __call__ (line 363) | def __call__(self, hidden_states, context, deterministic=True):
class FlaxFeedForward (line 388) | class FlaxFeedForward(nn.Module):
method setup (line 409) | def setup(self):
method __call__ (line 415) | def __call__(self, hidden_states, deterministic=True):
class FlaxGEGLU (line 421) | class FlaxGEGLU(nn.Module):
method setup (line 438) | def setup(self):
method __call__ (line 443) | def __call__(self, hidden_states, deterministic=True):
FILE: 4DoF/diffusers/models/attention_processor.py
function rotate_every_two (line 36) | def rotate_every_two(x):
function cape (line 42) | def cape(x, p):
function cape_embed (line 48) | def cape_embed(p1, p2, qq, kk):
class Attention (line 75) | class Attention(nn.Module):
method __init__ (line 90) | def __init__(
method set_use_memory_efficient_attention_xformers (line 205) | def set_use_memory_efficient_attention_xformers(
method set_attention_slice (line 323) | def set_attention_slice(self, slice_size):
method set_processor (line 344) | def set_processor(self, processor: "AttnProcessor"):
method forward (line 357) | def forward(self, hidden_states, encoder_hidden_states=None, attention...
method batch_to_head_dim (line 369) | def batch_to_head_dim(self, tensor):
method head_to_batch_dim (line 376) | def head_to_batch_dim(self, tensor, out_dim=3):
method get_attention_scores (line 387) | def get_attention_scores(self, query, key, attention_mask=None):
method prepare_attention_mask (line 421) | def prepare_attention_mask(self, attention_mask, target_length, batch_...
method norm_encoder_hidden_states (line 462) | def norm_encoder_hidden_states(self, encoder_hidden_states):
class AttnProcessor (line 482) | class AttnProcessor:
method __call__ (line 487) | def __call__(
class LoRALinearLayer (line 548) | class LoRALinearLayer(nn.Module):
method __init__ (line 549) | def __init__(self, in_features, out_features, rank=4, network_alpha=No...
method forward (line 565) | def forward(self, hidden_states):
class LoRAAttnProcessor (line 578) | class LoRAAttnProcessor(nn.Module):
method __init__ (line 593) | def __init__(self, hidden_size, cross_attention_dim=None, rank=4, netw...
method __call__ (line 605) | def __call__(
class CustomDiffusionAttnProcessor (line 661) | class CustomDiffusionAttnProcessor(nn.Module):
method __init__ (line 680) | def __init__(
method __call__ (line 706) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class AttnAddedKVProcessor (line 757) | class AttnAddedKVProcessor:
method __call__ (line 763) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class AttnAddedKVProcessor2_0 (line 811) | class AttnAddedKVProcessor2_0:
method __init__ (line 817) | def __init__(self):
method __call__ (line 823) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class LoRAAttnAddedKVProcessor (line 874) | class LoRAAttnAddedKVProcessor(nn.Module):
method __init__ (line 889) | def __init__(self, hidden_size, cross_attention_dim=None, rank=4, netw...
method __call__ (line 903) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class XFormersAttnAddedKVProcessor (line 955) | class XFormersAttnAddedKVProcessor:
method __init__ (line 967) | def __init__(self, attention_op: Optional[Callable] = None):
method __call__ (line 970) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class XFormersAttnProcessor (line 1020) | class XFormersAttnProcessor:
method __init__ (line 1032) | def __init__(self, attention_op: Optional[Callable] = None):
method __call__ (line 1035) | def __call__(
class AttnProcessor2_0 (line 1135) | class AttnProcessor2_0:
method __init__ (line 1140) | def __init__(self):
method __call__ (line 1144) | def __call__(
class LoRAXFormersAttnProcessor (line 1219) | class LoRAXFormersAttnProcessor(nn.Module):
method __init__ (line 1240) | def __init__(
method __call__ (line 1255) | def __call__(
class LoRAAttnProcessor2_0 (line 1312) | class LoRAAttnProcessor2_0(nn.Module):
method __init__ (line 1328) | def __init__(self, hidden_size, cross_attention_dim=None, rank=4, netw...
method __call__ (line 1342) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class CustomDiffusionXFormersAttnProcessor (line 1403) | class CustomDiffusionXFormersAttnProcessor(nn.Module):
method __init__ (line 1426) | def __init__(
method __call__ (line 1454) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class SlicedAttnProcessor (line 1510) | class SlicedAttnProcessor:
method __init__ (line 1520) | def __init__(self, slice_size):
method __call__ (line 1523) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class SlicedAttnAddedKVProcessor (line 1591) | class SlicedAttnAddedKVProcessor:
method __init__ (line 1601) | def __init__(self, slice_size):
method __call__ (line 1604) | def __call__(self, attn: "Attention", hidden_states, encoder_hidden_st...
class SpatialNorm (line 1694) | class SpatialNorm(nn.Module):
method __init__ (line 1699) | def __init__(
method forward (line 1709) | def forward(self, f, zq):
FILE: 4DoF/diffusers/models/autoencoder_kl.py
class AutoencoderKLOutput (line 28) | class AutoencoderKLOutput(BaseOutput):
class AutoencoderKL (line 41) | class AutoencoderKL(ModelMixin, ConfigMixin):
method __init__ (line 72) | def __init__(
method _set_gradient_checkpointing (line 127) | def _set_gradient_checkpointing(self, module, value=False):
method enable_tiling (line 131) | def enable_tiling(self, use_tiling: bool = True):
method disable_tiling (line 139) | def disable_tiling(self):
method enable_slicing (line 146) | def enable_slicing(self):
method disable_slicing (line 153) | def disable_slicing(self):
method attn_processors (line 162) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 186) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 221) | def set_default_attn_processor(self):
method encode (line 228) | def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> Au...
method _decode (line 246) | def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> U...
method decode (line 259) | def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Un...
method blend_v (line 271) | def blend_v(self, a, b, blend_extent):
method blend_h (line 277) | def blend_h(self, a, b, blend_extent):
method tiled_encode (line 283) | def tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True)...
method tiled_decode (line 337) | def tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True)...
method forward (line 385) | def forward(
FILE: 4DoF/diffusers/models/controlnet.py
class ControlNetOutput (line 39) | class ControlNetOutput(BaseOutput):
class ControlNetConditioningEmbedding (line 58) | class ControlNetConditioningEmbedding(nn.Module):
method __init__ (line 68) | def __init__(
method forward (line 90) | def forward(self, conditioning):
class ControlNetModel (line 103) | class ControlNetModel(ModelMixin, ConfigMixin):
method __init__ (line 159) | def __init__(
method from_unet (line 344) | def from_unet(
method attn_processors (line 400) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 424) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 459) | def set_default_attn_processor(self):
method set_attention_slice (line 466) | def set_attention_slice(self, slice_size):
method _set_gradient_checkpointing (line 531) | def _set_gradient_checkpointing(self, module, value=False):
method forward (line 535) | def forward(
function zero_module (line 702) | def zero_module(module):
FILE: 4DoF/diffusers/models/controlnet_flax.py
class FlaxControlNetOutput (line 34) | class FlaxControlNetOutput(BaseOutput):
class FlaxControlNetConditioningEmbedding (line 47) | class FlaxControlNetConditioningEmbedding(nn.Module):
method setup (line 52) | def setup(self):
method __call__ (line 90) | def __call__(self, conditioning):
class FlaxControlNetModel (line 104) | class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin):
method init_weights (line 171) | def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict:
method setup (line 185) | def setup(self):
method __call__ (line 313) | def __call__(
FILE: 4DoF/diffusers/models/cross_attention.py
class CrossAttention (line 41) | class CrossAttention(Attention):
method __init__ (line 42) | def __init__(self, *args, **kwargs):
class CrossAttnProcessor (line 48) | class CrossAttnProcessor(AttnProcessorRename):
method __init__ (line 49) | def __init__(self, *args, **kwargs):
class LoRACrossAttnProcessor (line 55) | class LoRACrossAttnProcessor(LoRAAttnProcessor):
method __init__ (line 56) | def __init__(self, *args, **kwargs):
class CrossAttnAddedKVProcessor (line 62) | class CrossAttnAddedKVProcessor(AttnAddedKVProcessor):
method __init__ (line 63) | def __init__(self, *args, **kwargs):
class XFormersCrossAttnProcessor (line 69) | class XFormersCrossAttnProcessor(XFormersAttnProcessor):
method __init__ (line 70) | def __init__(self, *args, **kwargs):
class LoRAXFormersCrossAttnProcessor (line 76) | class LoRAXFormersCrossAttnProcessor(LoRAXFormersAttnProcessor):
method __init__ (line 77) | def __init__(self, *args, **kwargs):
class SlicedCrossAttnProcessor (line 83) | class SlicedCrossAttnProcessor(SlicedAttnProcessor):
method __init__ (line 84) | def __init__(self, *args, **kwargs):
class SlicedCrossAttnAddedKVProcessor (line 90) | class SlicedCrossAttnAddedKVProcessor(SlicedAttnAddedKVProcessor):
method __init__ (line 91) | def __init__(self, *args, **kwargs):
FILE: 4DoF/diffusers/models/dual_transformer_2d.py
class DualTransformer2DModel (line 21) | class DualTransformer2DModel(nn.Module):
method __init__ (line 48) | def __init__(
method forward (line 97) | def forward(
FILE: 4DoF/diffusers/models/embeddings.py
function get_timestep_embedding (line 24) | def get_timestep_embedding(
function get_2d_sincos_pos_embed (line 67) | def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra...
function get_2d_sincos_pos_embed_from_grid (line 84) | def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
function get_1d_sincos_pos_embed_from_grid (line 96) | def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
class PatchEmbed (line 117) | class PatchEmbed(nn.Module):
method __init__ (line 120) | def __init__(
method forward (line 148) | def forward(self, latent):
class TimestepEmbedding (line 157) | class TimestepEmbedding(nn.Module):
method __init__ (line 158) | def __init__(
method forward (line 189) | def forward(self, sample, condition=None):
class Timesteps (line 204) | class Timesteps(nn.Module):
method __init__ (line 205) | def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale...
method forward (line 211) | def forward(self, timesteps):
class GaussianFourierProjection (line 221) | class GaussianFourierProjection(nn.Module):
method __init__ (line 224) | def __init__(
method forward (line 238) | def forward(self, x):
class ImagePositionalEmbeddings (line 251) | class ImagePositionalEmbeddings(nn.Module):
method __init__ (line 275) | def __init__(
method forward (line 293) | def forward(self, index):
class LabelEmbedding (line 316) | class LabelEmbedding(nn.Module):
method __init__ (line 326) | def __init__(self, num_classes, hidden_size, dropout_prob):
method token_drop (line 333) | def token_drop(self, labels, force_drop_ids=None):
method forward (line 344) | def forward(self, labels: torch.LongTensor, force_drop_ids=None):
class TextImageProjection (line 352) | class TextImageProjection(nn.Module):
method __init__ (line 353) | def __init__(
method forward (line 366) | def forward(self, text_embeds: torch.FloatTensor, image_embeds: torch....
class ImageProjection (line 379) | class ImageProjection(nn.Module):
method __init__ (line 380) | def __init__(
method forward (line 392) | def forward(self, image_embeds: torch.FloatTensor):
class CombinedTimestepLabelEmbeddings (line 402) | class CombinedTimestepLabelEmbeddings(nn.Module):
method __init__ (line 403) | def __init__(self, num_classes, embedding_dim, class_dropout_prob=0.1):
method forward (line 410) | def forward(self, timestep, class_labels, hidden_dtype=None):
class TextTimeEmbedding (line 421) | class TextTimeEmbedding(nn.Module):
method __init__ (line 422) | def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: i...
method forward (line 429) | def forward(self, hidden_states):
class TextImageTimeEmbedding (line 437) | class TextImageTimeEmbedding(nn.Module):
method __init__ (line 438) | def __init__(self, text_embed_dim: int = 768, image_embed_dim: int = 7...
method forward (line 444) | def forward(self, text_embeds: torch.FloatTensor, image_embeds: torch....
class ImageTimeEmbedding (line 455) | class ImageTimeEmbedding(nn.Module):
method __init__ (line 456) | def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1...
method forward (line 461) | def forward(self, image_embeds: torch.FloatTensor):
class ImageHintTimeEmbedding (line 468) | class ImageHintTimeEmbedding(nn.Module):
method __init__ (line 469) | def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1...
method forward (line 491) | def forward(self, image_embeds: torch.FloatTensor, hint: torch.FloatTe...
class AttentionPooling (line 499) | class AttentionPooling(nn.Module):
method __init__ (line 502) | def __init__(self, num_heads, embed_dim, dtype=None):
method forward (line 512) | def forward(self, x):
FILE: 4DoF/diffusers/models/embeddings_flax.py
function get_sinusoidal_embeddings (line 20) | def get_sinusoidal_embeddings(
class FlaxTimestepEmbedding (line 58) | class FlaxTimestepEmbedding(nn.Module):
method __call__ (line 72) | def __call__(self, temb):
class FlaxTimesteps (line 79) | class FlaxTimesteps(nn.Module):
method __call__ (line 92) | def __call__(self, timesteps):
FILE: 4DoF/diffusers/models/modeling_flax_pytorch_utils.py
function rename_key (line 28) | def rename_key(key):
function rename_key_and_reshape_tensor (line 43) | def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_s...
function convert_pytorch_state_dict_to_flax (line 90) | def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model, init_k...
FILE: 4DoF/diffusers/models/modeling_flax_utils.py
class FlaxModelMixin (line 45) | class FlaxModelMixin:
method _from_config (line 59) | def _from_config(cls, config, **kwargs):
method _cast_floating_to (line 65) | def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jn...
method to_bf16 (line 89) | def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
method to_fp32 (line 128) | def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):
method to_fp16 (line 155) | def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
method init_weights (line 194) | def init_weights(self, rng: jax.random.KeyArray) -> Dict:
method from_pretrained (line 198) | def from_pretrained(
method save_pretrained (line 495) | def save_pretrained(
FILE: 4DoF/diffusers/models/modeling_pytorch_flax_utils.py
function load_flax_checkpoint_in_pytorch_model (line 37) | def load_flax_checkpoint_in_pytorch_model(pt_model, model_file):
function load_flax_weights_in_pytorch_model (line 58) | def load_flax_weights_in_pytorch_model(pt_model, flax_state):
FILE: 4DoF/diffusers/models/modeling_utils.py
function get_parameter_device (line 63) | def get_parameter_device(parameter: torch.nn.Module):
function get_parameter_dtype (line 79) | def get_parameter_dtype(parameter: torch.nn.Module):
function load_state_dict (line 101) | def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: O...
function _load_state_dict_into_model (line 132) | def _load_state_dict_into_model(model_to_load, state_dict):
class ModelMixin (line 153) | class ModelMixin(torch.nn.Module):
method __init__ (line 167) | def __init__(self):
method __getattr__ (line 170) | def __getattr__(self, name: str) -> Any:
method is_gradient_checkpointing (line 189) | def is_gradient_checkpointing(self) -> bool:
method enable_gradient_checkpointing (line 195) | def enable_gradient_checkpointing(self):
method disable_gradient_checkpointing (line 204) | def disable_gradient_checkpointing(self):
method set_use_memory_efficient_attention_xformers (line 212) | def set_use_memory_efficient_attention_xformers(
method enable_xformers_memory_efficient_attention (line 229) | def enable_xformers_memory_efficient_attention(self, attention_op: Opt...
method disable_xformers_memory_efficient_attention (line 265) | def disable_xformers_memory_efficient_attention(self):
method save_pretrained (line 271) | def save_pretrained(
method from_pretrained (line 332) | def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union...
method _load_pretrained_model (line 732) | def _load_pretrained_model(
method device (line 836) | def device(self) -> device:
method dtype (line 844) | def dtype(self) -> torch.dtype:
method num_parameters (line 850) | def num_parameters(self, only_trainable: bool = False, exclude_embeddi...
method _convert_deprecated_attention_blocks (line 888) | def _convert_deprecated_attention_blocks(self, state_dict):
method _temp_convert_self_to_deprecated_attention_blocks (line 932) | def _temp_convert_self_to_deprecated_attention_blocks(self):
method _undo_temp_convert_self_to_deprecated_attention_blocks (line 959) | def _undo_temp_convert_self_to_deprecated_attention_blocks(self):
FILE: 4DoF/diffusers/models/prior_transformer.py
class PriorTransformerOutput (line 17) | class PriorTransformerOutput(BaseOutput):
class PriorTransformer (line 29) | class PriorTransformer(ModelMixin, ConfigMixin):
method __init__ (line 67) | def __init__(
method attn_processors (line 164) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 188) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 223) | def set_default_attn_processor(self):
method forward (line 229) | def forward(
method post_process_latents (line 362) | def post_process_latents(self, prior_latents):
FILE: 4DoF/diffusers/models/resnet.py
class Upsample1D (line 28) | class Upsample1D(nn.Module):
method __init__ (line 42) | def __init__(self, channels, use_conv=False, use_conv_transpose=False,...
method forward (line 56) | def forward(self, inputs):
class Downsample1D (line 69) | class Downsample1D(nn.Module):
method __init__ (line 83) | def __init__(self, channels, use_conv=False, out_channels=None, paddin...
method forward (line 98) | def forward(self, inputs):
class Upsample2D (line 103) | class Upsample2D(nn.Module):
method __init__ (line 117) | def __init__(self, channels, use_conv=False, use_conv_transpose=False,...
method forward (line 137) | def forward(self, hidden_states, output_size=None):
class Downsample2D (line 175) | class Downsample2D(nn.Module):
method __init__ (line 189) | def __init__(self, channels, use_conv=False, out_channels=None, paddin...
method forward (line 213) | def forward(self, hidden_states):
class FirUpsample2D (line 225) | class FirUpsample2D(nn.Module):
method __init__ (line 239) | def __init__(self, channels=None, out_channels=None, use_conv=False, f...
method _upsample_2d (line 248) | def _upsample_2d(self, hidden_states, weight=None, kernel=None, factor...
method forward (line 328) | def forward(self, hidden_states):
class FirDownsample2D (line 338) | class FirDownsample2D(nn.Module):
method __init__ (line 352) | def __init__(self, channels=None, out_channels=None, use_conv=False, f...
method _downsample_2d (line 361) | def _downsample_2d(self, hidden_states, weight=None, kernel=None, fact...
method forward (line 415) | def forward(self, hidden_states):
class KDownsample2D (line 426) | class KDownsample2D(nn.Module):
method __init__ (line 427) | def __init__(self, pad_mode="reflect"):
method forward (line 434) | def forward(self, inputs):
class KUpsample2D (line 443) | class KUpsample2D(nn.Module):
method __init__ (line 444) | def __init__(self, pad_mode="reflect"):
method forward (line 451) | def forward(self, inputs):
class ResnetBlock2D (line 460) | class ResnetBlock2D(nn.Module):
method __init__ (line 491) | def __init__(
method forward (line 590) | def forward(self, input_tensor, temb):
function rearrange_dims (line 644) | def rearrange_dims(tensor):
class Conv1dBlock (line 655) | class Conv1dBlock(nn.Module):
method __init__ (line 660) | def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):
method forward (line 667) | def forward(self, inputs):
class ResidualTemporalBlock1D (line 677) | class ResidualTemporalBlock1D(nn.Module):
method __init__ (line 678) | def __init__(self, inp_channels, out_channels, embed_dim, kernel_size=5):
method forward (line 690) | def forward(self, inputs, t):
function upsample_2d (line 706) | def upsample_2d(hidden_states, kernel=None, factor=2, gain=1):
function downsample_2d (line 743) | def downsample_2d(hidden_states, kernel=None, factor=2, gain=1):
function upfirdn2d_native (line 778) | def upfirdn2d_native(tensor, kernel, up=1, down=1, pad=(0, 0)):
class TemporalConvLayer (line 822) | class TemporalConvLayer(nn.Module):
method __init__ (line 828) | def __init__(self, in_dim, out_dim=None, dropout=0.0):
method forward (line 861) | def forward(self, hidden_states, num_frames=1):
FILE: 4DoF/diffusers/models/resnet_flax.py
class FlaxUpsample2D (line 19) | class FlaxUpsample2D(nn.Module):
method setup (line 23) | def setup(self):
method __call__ (line 32) | def __call__(self, hidden_states):
class FlaxDownsample2D (line 43) | class FlaxDownsample2D(nn.Module):
method setup (line 47) | def setup(self):
method __call__ (line 56) | def __call__(self, hidden_states):
class FlaxResnetBlock2D (line 63) | class FlaxResnetBlock2D(nn.Module):
method setup (line 70) | def setup(self):
method __call__ (line 106) | def __call__(self, hidden_states, temb, deterministic=True):
FILE: 4DoF/diffusers/models/t5_film_transformer.py
class T5FilmDecoder (line 25) | class T5FilmDecoder(ModelMixin, ConfigMixin):
method __init__ (line 27) | def __init__(
method encoder_decoder_mask (line 66) | def encoder_decoder_mask(self, query_input, key_input):
method forward (line 70) | def forward(self, encodings_and_masks, decoder_input_tokens, decoder_n...
class DecoderLayer (line 127) | class DecoderLayer(nn.Module):
method __init__ (line 128) | def __init__(self, d_model, d_kv, num_heads, d_ff, dropout_rate, layer...
method forward (line 153) | def forward(
class T5LayerSelfAttentionCond (line 185) | class T5LayerSelfAttentionCond(nn.Module):
method __init__ (line 186) | def __init__(self, d_model, d_kv, num_heads, dropout_rate):
method forward (line 193) | def forward(
class T5LayerCrossAttention (line 213) | class T5LayerCrossAttention(nn.Module):
method __init__ (line 214) | def __init__(self, d_model, d_kv, num_heads, dropout_rate, layer_norm_...
method forward (line 220) | def forward(
class T5LayerFFCond (line 236) | class T5LayerFFCond(nn.Module):
method __init__ (line 237) | def __init__(self, d_model, d_ff, dropout_rate, layer_norm_epsilon):
method forward (line 244) | def forward(self, hidden_states, conditioning_emb=None):
class T5DenseGatedActDense (line 254) | class T5DenseGatedActDense(nn.Module):
method __init__ (line 255) | def __init__(self, d_model, d_ff, dropout_rate):
method forward (line 263) | def forward(self, hidden_states):
class T5LayerNorm (line 273) | class T5LayerNorm(nn.Module):
method __init__ (line 274) | def __init__(self, hidden_size, eps=1e-6):
method forward (line 282) | def forward(self, hidden_states):
class NewGELUActivation (line 298) | class NewGELUActivation(nn.Module):
method forward (line 304) | def forward(self, input: torch.Tensor) -> torch.Tensor:
class T5FiLMLayer (line 308) | class T5FiLMLayer(nn.Module):
method __init__ (line 313) | def __init__(self, in_features, out_features):
method forward (line 317) | def forward(self, x, conditioning_emb):
FILE: 4DoF/diffusers/models/transformer_2d.py
class Transformer2DModelOutput (line 30) | class Transformer2DModelOutput(BaseOutput):
class Transformer2DModel (line 43) | class Transformer2DModel(ModelMixin, ConfigMixin):
method __init__ (line 72) | def __init__(
method forward (line 206) | def forward(
FILE: 4DoF/diffusers/models/transformer_temporal.py
class TransformerTemporalModelOutput (line 27) | class TransformerTemporalModelOutput(BaseOutput):
class TransformerTemporalModel (line 39) | class TransformerTemporalModel(ModelMixin, ConfigMixin):
method __init__ (line 61) | def __init__(
method forward (line 107) | def forward(
FILE: 4DoF/diffusers/models/unet_1d.py
class UNet1DOutput (line 29) | class UNet1DOutput(BaseOutput):
class UNet1DModel (line 41) | class UNet1DModel(ModelMixin, ConfigMixin):
method __init__ (line 75) | def __init__(
method forward (line 195) | def forward(
FILE: 4DoF/diffusers/models/unet_1d_blocks.py
class DownResnetBlock1D (line 24) | class DownResnetBlock1D(nn.Module):
method __init__ (line 25) | def __init__(
method forward (line 68) | def forward(self, hidden_states, temb=None):
class UpResnetBlock1D (line 86) | class UpResnetBlock1D(nn.Module):
method __init__ (line 87) | def __init__(
method forward (line 128) | def forward(self, hidden_states, res_hidden_states_tuple=None, temb=No...
class ValueFunctionMidBlock1D (line 146) | class ValueFunctionMidBlock1D(nn.Module):
method __init__ (line 147) | def __init__(self, in_channels, out_channels, embed_dim):
method forward (line 158) | def forward(self, x, temb=None):
class MidResTemporalBlock1D (line 166) | class MidResTemporalBlock1D(nn.Module):
method __init__ (line 167) | def __init__(
method forward (line 206) | def forward(self, hidden_states, temb):
class OutConv1DBlock (line 219) | class OutConv1DBlock(nn.Module):
method __init__ (line 220) | def __init__(self, num_groups_out, out_channels, embed_dim, act_fn):
method forward (line 227) | def forward(self, hidden_states, temb=None):
class OutValueFunctionBlock (line 237) | class OutValueFunctionBlock(nn.Module):
method __init__ (line 238) | def __init__(self, fc_dim, embed_dim):
method forward (line 248) | def forward(self, hidden_states, temb):
class Downsample1d (line 277) | class Downsample1d(nn.Module):
method __init__ (line 278) | def __init__(self, kernel="linear", pad_mode="reflect"):
method forward (line 285) | def forward(self, hidden_states):
class Upsample1d (line 294) | class Upsample1d(nn.Module):
method __init__ (line 295) | def __init__(self, kernel="linear", pad_mode="reflect"):
method forward (line 302) | def forward(self, hidden_states, temb=None):
class SelfAttention1d (line 311) | class SelfAttention1d(nn.Module):
method __init__ (line 312) | def __init__(self, in_channels, n_head=1, dropout_rate=0.0):
method transpose_for_scores (line 326) | def transpose_for_scores(self, projection: torch.Tensor) -> torch.Tensor:
method forward (line 332) | def forward(self, hidden_states):
class ResConvBlock (line 369) | class ResConvBlock(nn.Module):
method __init__ (line 370) | def __init__(self, in_channels, mid_channels, out_channels, is_last=Fa...
method forward (line 387) | def forward(self, hidden_states):
class UNetMidBlock1D (line 403) | class UNetMidBlock1D(nn.Module):
method __init__ (line 404) | def __init__(self, mid_channels, in_channels, out_channels=None):
method forward (line 432) | def forward(self, hidden_states, temb=None):
class AttnDownBlock1D (line 443) | class AttnDownBlock1D(nn.Module):
method __init__ (line 444) | def __init__(self, out_channels, in_channels, mid_channels=None):
method forward (line 463) | def forward(self, hidden_states, temb=None):
class DownBlock1D (line 473) | class DownBlock1D(nn.Module):
method __init__ (line 474) | def __init__(self, out_channels, in_channels, mid_channels=None):
method forward (line 487) | def forward(self, hidden_states, temb=None):
class DownBlock1DNoSkip (line 496) | class DownBlock1DNoSkip(nn.Module):
method __init__ (line 497) | def __init__(self, out_channels, in_channels, mid_channels=None):
method forward (line 509) | def forward(self, hidden_states, temb=None):
class AttnUpBlock1D (line 517) | class AttnUpBlock1D(nn.Module):
method __init__ (line 518) | def __init__(self, in_channels, out_channels, mid_channels=None):
method forward (line 537) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
class UpBlock1D (line 550) | class UpBlock1D(nn.Module):
method __init__ (line 551) | def __init__(self, in_channels, out_channels, mid_channels=None):
method forward (line 564) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
class UpBlock1DNoSkip (line 576) | class UpBlock1DNoSkip(nn.Module):
method __init__ (line 577) | def __init__(self, in_channels, out_channels, mid_channels=None):
method forward (line 589) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
function get_down_block (line 599) | def get_down_block(down_block_type, num_layers, in_channels, out_channel...
function get_up_block (line 617) | def get_up_block(up_block_type, num_layers, in_channels, out_channels, t...
function get_mid_block (line 635) | def get_mid_block(mid_block_type, num_layers, in_channels, mid_channels,...
function get_out_block (line 651) | def get_out_block(*, out_block_type, num_groups_out, embed_dim, out_chan...
FILE: 4DoF/diffusers/models/unet_2d.py
class UNet2DOutput (line 28) | class UNet2DOutput(BaseOutput):
class UNet2DModel (line 40) | class UNet2DModel(ModelMixin, ConfigMixin):
method __init__ (line 88) | def __init__(
method forward (line 228) | def forward(
FILE: 4DoF/diffusers/models/unet_2d_blocks.py
function get_down_block (line 32) | def get_down_block(
function get_up_block (line 227) | def get_up_block(
class UNetMidBlock2D (line 426) | class UNetMidBlock2D(nn.Module):
method __init__ (line 427) | def __init__(
method forward (line 507) | def forward(self, hidden_states, temb=None):
class UNetMidBlock2DCrossAttn (line 517) | class UNetMidBlock2DCrossAttn(nn.Module):
method __init__ (line 518) | def __init__(
method forward (line 603) | def forward(
class UNetMidBlock2DSimpleCrossAttn (line 629) | class UNetMidBlock2DSimpleCrossAttn(nn.Module):
method __init__ (line 630) | def __init__(
method forward (line 714) | def forward(
class AttnDownBlock2D (line 752) | class AttnDownBlock2D(nn.Module):
method __init__ (line 753) | def __init__(
method forward (line 844) | def forward(self, hidden_states, temb=None, upsample_size=None):
class CrossAttnDownBlock2D (line 864) | class CrossAttnDownBlock2D(nn.Module):
method __init__ (line 865) | def __init__(
method forward (line 952) | def forward(
class DownBlock2D (line 1018) | class DownBlock2D(nn.Module):
method __init__ (line 1019) | def __init__(
method forward (line 1070) | def forward(self, hidden_states, temb=None):
class DownEncoderBlock2D (line 1104) | class DownEncoderBlock2D(nn.Module):
method __init__ (line 1105) | def __init__(
method forward (line 1153) | def forward(self, hidden_states):
class AttnDownEncoderBlock2D (line 1164) | class AttnDownEncoderBlock2D(nn.Module):
method __init__ (line 1165) | def __init__(
method forward (line 1236) | def forward(self, hidden_states):
class AttnSkipDownBlock2D (line 1248) | class AttnSkipDownBlock2D(nn.Module):
method __init__ (line 1249) | def __init__(
method forward (line 1329) | def forward(self, hidden_states, temb=None, skip_sample=None):
class SkipDownBlock2D (line 1349) | class SkipDownBlock2D(nn.Module):
method __init__ (line 1350) | def __init__(
method forward (line 1409) | def forward(self, hidden_states, temb=None, skip_sample=None):
class ResnetDownsampleBlock2D (line 1428) | class ResnetDownsampleBlock2D(nn.Module):
method __init__ (line 1429) | def __init__(
method forward (line 1492) | def forward(self, hidden_states, temb=None):
class SimpleCrossAttnDownBlock2D (line 1526) | class SimpleCrossAttnDownBlock2D(nn.Module):
method __init__ (line 1527) | def __init__(
method forward (line 1621) | def forward(
class KDownBlock2D (line 1685) | class KDownBlock2D(nn.Module):
method __init__ (line 1686) | def __init__(
method forward (line 1731) | def forward(self, hidden_states, temb=None):
class KCrossAttnDownBlock2D (line 1763) | class KCrossAttnDownBlock2D(nn.Module):
method __init__ (line 1764) | def __init__(
method forward (line 1828) | def forward(
class AttnUpBlock2D (line 1891) | class AttnUpBlock2D(nn.Module):
method __init__ (line 1892) | def __init__(
method forward (line 1980) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
class CrossAttnUpBlock2D (line 2000) | class CrossAttnUpBlock2D(nn.Module):
method __init__ (line 2001) | def __init__(
method forward (line 2084) | def forward(
class UpBlock2D (line 2151) | class UpBlock2D(nn.Module):
method __init__ (line 2152) | def __init__(
method forward (line 2199) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
class UpDecoderBlock2D (line 2232) | class UpDecoderBlock2D(nn.Module):
method __init__ (line 2233) | def __init__(
method forward (line 2276) | def forward(self, hidden_states, temb=None):
class AttnUpDecoderBlock2D (line 2287) | class AttnUpDecoderBlock2D(nn.Module):
method __init__ (line 2288) | def __init__(
method forward (line 2355) | def forward(self, hidden_states, temb=None):
class AttnSkipUpBlock2D (line 2367) | class AttnSkipUpBlock2D(nn.Module):
method __init__ (line 2368) | def __init__(
method forward (line 2458) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, s...
class SkipUpBlock2D (line 2486) | class SkipUpBlock2D(nn.Module):
method __init__ (line 2487) | def __init__(
method forward (line 2555) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, s...
class ResnetUpsampleBlock2D (line 2581) | class ResnetUpsampleBlock2D(nn.Module):
method __init__ (line 2582) | def __init__(
method forward (line 2648) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
class SimpleCrossAttnUpBlock2D (line 2681) | class SimpleCrossAttnUpBlock2D(nn.Module):
method __init__ (line 2682) | def __init__(
method forward (line 2778) | def forward(
class KUpBlock2D (line 2845) | class KUpBlock2D(nn.Module):
method __init__ (line 2846) | def __init__(
method forward (line 2893) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
class KCrossAttnUpBlock2D (line 2925) | class KCrossAttnUpBlock2D(nn.Module):
method __init__ (line 2926) | def __init__(
method forward (line 3009) | def forward(
class KAttentionBlock (line 3072) | class KAttentionBlock(nn.Module):
method __init__ (line 3089) | def __init__(
method _to_3d (line 3132) | def _to_3d(self, hidden_states, height, weight):
method _to_4d (line 3135) | def _to_4d(self, hidden_states, height, weight):
method forward (line 3138) | def forward(
FILE: 4DoF/diffusers/models/unet_2d_blocks_flax.py
class FlaxCrossAttnDownBlock2D (line 22) | class FlaxCrossAttnDownBlock2D(nn.Module):
method setup (line 56) | def setup(self):
method __call__ (line 89) | def __call__(self, hidden_states, temb, encoder_hidden_states, determi...
class FlaxDownBlock2D (line 104) | class FlaxDownBlock2D(nn.Module):
method setup (line 129) | def setup(self):
method __call__ (line 147) | def __call__(self, hidden_states, temb, deterministic=True):
class FlaxCrossAttnUpBlock2D (line 161) | class FlaxCrossAttnUpBlock2D(nn.Module):
method setup (line 196) | def setup(self):
method __call__ (line 230) | def __call__(self, hidden_states, res_hidden_states_tuple, temb, encod...
class FlaxUpBlock2D (line 246) | class FlaxUpBlock2D(nn.Module):
method setup (line 274) | def setup(self):
method __call__ (line 294) | def __call__(self, hidden_states, res_hidden_states_tuple, temb, deter...
class FlaxUNetMidBlock2DCrossAttn (line 309) | class FlaxUNetMidBlock2DCrossAttn(nn.Module):
method setup (line 335) | def setup(self):
method __call__ (line 371) | def __call__(self, hidden_states, temb, encoder_hidden_states, determi...
FILE: 4DoF/diffusers/models/unet_2d_condition.py
class UNet2DConditionOutput (line 54) | class UNet2DConditionOutput(BaseOutput):
class UNet2DConditionModel (line 66) | class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoade...
method __init__ (line 157) | def __init__(
method attn_processors (line 568) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 591) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 625) | def set_default_attn_processor(self):
method set_attention_slice (line 631) | def set_attention_slice(self, slice_size):
method _set_gradient_checkpointing (line 696) | def _set_gradient_checkpointing(self, module, value=False):
method forward (line 700) | def forward(
FILE: 4DoF/diffusers/models/unet_2d_condition_flax.py
class FlaxUNet2DConditionOutput (line 36) | class FlaxUNet2DConditionOutput(BaseOutput):
class FlaxUNet2DConditionModel (line 49) | class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
method init_weights (line 120) | def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict:
method setup (line 132) | def setup(self):
method __call__ (line 267) | def __call__(
FILE: 4DoF/diffusers/models/unet_3d_blocks.py
function get_down_block (line 23) | def get_down_block(
function get_up_block (line 79) | def get_up_block(
class UNetMidBlock3DCrossAttn (line 135) | class UNetMidBlock3DCrossAttn(nn.Module):
method __init__ (line 136) | def __init__(
method forward (line 235) | def forward(
class CrossAttnDownBlock3D (line 264) | class CrossAttnDownBlock3D(nn.Module):
method __init__ (line 265) | def __init__(
method forward (line 360) | def forward(
class DownBlock3D (line 398) | class DownBlock3D(nn.Module):
method __init__ (line 399) | def __init__(
method forward (line 459) | def forward(self, hidden_states, temb=None, num_frames=1):
class CrossAttnUpBlock3D (line 477) | class CrossAttnUpBlock3D(nn.Module):
method __init__ (line 478) | def __init__(
method forward (line 569) | def forward(
class UpBlock3D (line 608) | class UpBlock3D(nn.Module):
method __init__ (line 609) | def __init__(
method forward (line 665) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
FILE: 4DoF/diffusers/models/unet_3d_condition.py
class UNet3DConditionOutput (line 44) | class UNet3DConditionOutput(BaseOutput):
class UNet3DConditionModel (line 56) | class UNet3DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoade...
method __init__ (line 90) | def __init__(
method attn_processors (line 273) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attention_slice (line 297) | def set_attention_slice(self, slice_size):
method set_attn_processor (line 363) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method enable_forward_chunking (line 397) | def enable_forward_chunking(self, chunk_size=None, dim=0):
method disable_forward_chunking (line 426) | def disable_forward_chunking(self):
method set_default_attn_processor (line 438) | def set_default_attn_processor(self):
method _set_gradient_checkpointing (line 444) | def _set_gradient_checkpointing(self, module, value=False):
method forward (line 448) | def forward(
FILE: 4DoF/diffusers/models/vae.py
class DecoderOutput (line 27) | class DecoderOutput(BaseOutput):
class Encoder (line 39) | class Encoder(nn.Module):
method __init__ (line 40) | def __init__(
method forward (line 108) | def forward(self, x):
class Decoder (line 152) | class Decoder(nn.Module):
method __init__ (line 153) | def __init__(
method forward (line 228) | def forward(self, z, latent_embeds=None):
class VectorQuantizer (line 283) | class VectorQuantizer(nn.Module):
method __init__ (line 292) | def __init__(
method remap_to_used (line 321) | def remap_to_used(self, inds):
method unmap_to_all (line 335) | def unmap_to_all(self, inds):
method forward (line 345) | def forward(self, z):
method get_codebook_entry (line 379) | def get_codebook_entry(self, indices, shape):
class DiagonalGaussianDistribution (line 397) | class DiagonalGaussianDistribution(object):
method __init__ (line 398) | def __init__(self, parameters, deterministic=False):
method sample (line 410) | def sample(self, generator: Optional[torch.Generator] = None) -> torch...
method kl (line 418) | def kl(self, other=None):
method nll (line 434) | def nll(self, sample, dims=[1, 2, 3]):
method mode (line 440) | def mode(self):
FILE: 4DoF/diffusers/models/vae_flax.py
class FlaxDecoderOutput (line 33) | class FlaxDecoderOutput(BaseOutput):
class FlaxAutoencoderKLOutput (line 48) | class FlaxAutoencoderKLOutput(BaseOutput):
class FlaxUpsample2D (line 61) | class FlaxUpsample2D(nn.Module):
method setup (line 75) | def setup(self):
method __call__ (line 84) | def __call__(self, hidden_states):
class FlaxDownsample2D (line 95) | class FlaxDownsample2D(nn.Module):
method setup (line 109) | def setup(self):
method __call__ (line 118) | def __call__(self, hidden_states):
class FlaxResnetBlock2D (line 125) | class FlaxResnetBlock2D(nn.Module):
method setup (line 151) | def setup(self):
method __call__ (line 185) | def __call__(self, hidden_states, deterministic=True):
class FlaxAttentionBlock (line 202) | class FlaxAttentionBlock(nn.Module):
method setup (line 222) | def setup(self):
method transpose_for_scores (line 231) | def transpose_for_scores(self, projection):
method __call__ (line 239) | def __call__(self, hidden_states):
class FlaxDownEncoderBlock2D (line 274) | class FlaxDownEncoderBlock2D(nn.Module):
method setup (line 302) | def setup(self):
method __call__ (line 320) | def __call__(self, hidden_states, deterministic=True):
class FlaxUpDecoderBlock2D (line 330) | class FlaxUpDecoderBlock2D(nn.Module):
method setup (line 358) | def setup(self):
method __call__ (line 376) | def __call__(self, hidden_states, deterministic=True):
class FlaxUNetMidBlock2D (line 386) | class FlaxUNetMidBlock2D(nn.Module):
method setup (line 411) | def setup(self):
method __call__ (line 448) | def __call__(self, hidden_states, deterministic=True):
class FlaxEncoder (line 457) | class FlaxEncoder(nn.Module):
method setup (line 501) | def setup(self):
method __call__ (line 550) | def __call__(self, sample, deterministic: bool = True):
class FlaxDecoder (line 569) | class FlaxDecoder(nn.Module):
method setup (line 612) | def setup(self):
method __call__ (line 665) | def __call__(self, sample, deterministic: bool = True):
class FlaxDiagonalGaussianDistribution (line 683) | class FlaxDiagonalGaussianDistribution(object):
method __init__ (line 684) | def __init__(self, parameters, deterministic=False):
method sample (line 694) | def sample(self, key):
method kl (line 697) | def kl(self, other=None):
method nll (line 709) | def nll(self, sample, axis=[1, 2, 3]):
method mode (line 716) | def mode(self):
class FlaxAutoencoderKL (line 721) | class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):
method setup (line 783) | def setup(self):
method init_weights (line 820) | def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict:
method encode (line 830) | def encode(self, sample, deterministic: bool = True, return_dict: bool...
method decode (line 842) | def decode(self, latents, deterministic: bool = True, return_dict: boo...
method __call__ (line 856) | def __call__(self, sample, sample_posterior=False, deterministic: bool...
FILE: 4DoF/diffusers/models/vq_model.py
class VQEncoderOutput (line 27) | class VQEncoderOutput(BaseOutput):
class VQModel (line 39) | class VQModel(ModelMixin, ConfigMixin):
method __init__ (line 70) | def __init__(
method encode (line 120) | def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> VQ...
method decode (line 130) | def decode(
method forward (line 146) | def forward(self, sample: torch.FloatTensor, return_dict: bool = True)...
FILE: 4DoF/diffusers/optimization.py
class SchedulerType (line 30) | class SchedulerType(Enum):
function get_constant_schedule (line 40) | def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
function get_constant_schedule_with_warmup (line 56) | def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_s...
function get_piecewise_constant_schedule (line 81) | def get_piecewise_constant_schedule(optimizer: Optimizer, step_rules: st...
function get_linear_schedule_with_warmup (line 123) | def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_tra...
function get_cosine_schedule_with_warmup (line 152) | def get_cosine_schedule_with_warmup(
function get_cosine_with_hard_restarts_schedule_with_warmup (line 186) | def get_cosine_with_hard_restarts_schedule_with_warmup(
function get_polynomial_decay_schedule_with_warmup (line 221) | def get_polynomial_decay_schedule_with_warmup(
function get_scheduler (line 282) | def get_scheduler(
FILE: 4DoF/diffusers/pipelines/alt_diffusion/__init__.py
class AltDiffusionPipelineOutput (line 13) | class AltDiffusionPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/alt_diffusion/modeling_roberta_series.py
class TransformationModelOutput (line 11) | class TransformationModelOutput(ModelOutput):
class RobertaSeriesConfig (line 39) | class RobertaSeriesConfig(XLMRobertaConfig):
method __init__ (line 40) | def __init__(
class RobertaSeriesModelWithTransformation (line 58) | class RobertaSeriesModelWithTransformation(RobertaPreTrainedModel):
method __init__ (line 64) | def __init__(self, config):
method forward (line 74) | def forward(
FILE: 4DoF/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py
function rescale_noise_cfg (line 55) | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
class AltDiffusionPipeline (line 70) | class AltDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixi...
method __init__ (line 107) | def __init__(
method enable_vae_slicing (line 197) | def enable_vae_slicing(self):
method disable_vae_slicing (line 206) | def disable_vae_slicing(self):
method enable_vae_tiling (line 213) | def enable_vae_tiling(self):
method disable_vae_tiling (line 222) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 229) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 254) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 283) | def _execution_device(self):
method _encode_prompt (line 300) | def _encode_prompt(
method run_safety_checker (line 454) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 468) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 483) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 500) | def check_inputs(
method prepare_latents (line 547) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 566) | def __call__(
FILE: 4DoF/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py
function preprocess (line 71) | def preprocess(image):
class AltDiffusionImg2ImgPipeline (line 98) | class AltDiffusionImg2ImgPipeline(
method __init__ (line 137) | def __init__(
method enable_sequential_cpu_offload (line 227) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 252) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 281) | def _execution_device(self):
method _encode_prompt (line 298) | def _encode_prompt(
method run_safety_checker (line 452) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 466) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 481) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 498) | def check_inputs(
method get_timesteps (line 538) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 547) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method __call__ (line 606) | def __call__(
FILE: 4DoF/diffusers/pipelines/audio_diffusion/mel.py
class Mel (line 37) | class Mel(ConfigMixin, SchedulerMixin):
method __init__ (line 52) | def __init__(
method set_resolution (line 73) | def set_resolution(self, x_res: int, y_res: int):
method load_audio (line 85) | def load_audio(self, audio_file: str = None, raw_audio: np.ndarray = N...
method get_number_of_slices (line 101) | def get_number_of_slices(self) -> int:
method get_audio_slice (line 109) | def get_audio_slice(self, slice: int = 0) -> np.ndarray:
method get_sample_rate (line 120) | def get_sample_rate(self) -> int:
method audio_slice_to_image (line 128) | def audio_slice_to_image(self, slice: int) -> Image.Image:
method image_to_audio (line 145) | def image_to_audio(self, image: Image.Image) -> np.ndarray:
FILE: 4DoF/diffusers/pipelines/audio_diffusion/pipeline_audio_diffusion.py
class AudioDiffusionPipeline (line 30) | class AudioDiffusionPipeline(DiffusionPipeline):
method __init__ (line 44) | def __init__(
method get_default_steps (line 54) | def get_default_steps(self) -> int:
method __call__ (line 63) | def __call__(
method encode (line 199) | def encode(self, images: List[Image.Image], steps: int = 50) -> np.nda...
method slerp (line 236) | def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.T...
FILE: 4DoF/diffusers/pipelines/audioldm/pipeline_audioldm.py
class AudioLDMPipeline (line 46) | class AudioLDMPipeline(DiffusionPipeline):
method __init__ (line 72) | def __init__(
method enable_vae_slicing (line 94) | def enable_vae_slicing(self):
method disable_vae_slicing (line 104) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 111) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 129) | def _execution_device(self):
method _encode_prompt (line 146) | def _encode_prompt(
method decode_latents (line 285) | def decode_latents(self, latents):
method mel_spectrogram_to_waveform (line 290) | def mel_spectrogram_to_waveform(self, mel_spectrogram):
method prepare_extra_step_kwargs (line 300) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 317) | def check_inputs(
method prepare_latents (line 376) | def prepare_latents(self, batch_size, num_channels_latents, height, dt...
method __call__ (line 400) | def __call__(
FILE: 4DoF/diffusers/pipelines/consistency_models/pipeline_consistency_models.py
class ConsistencyModelPipeline (line 51) | class ConsistencyModelPipeline(DiffusionPipeline):
method __init__ (line 69) | def __init__(self, unet: UNet2DModel, scheduler: CMStochasticIterative...
method enable_sequential_cpu_offload (line 79) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 104) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 134) | def _execution_device(self):
method prepare_latents (line 151) | def prepare_latents(self, batch_size, num_channels, height, width, dty...
method postprocess_image (line 169) | def postprocess_image(self, sample: torch.FloatTensor, output_type: st...
method prepare_class_labels (line 189) | def prepare_class_labels(self, batch_size, device, class_labels=None):
method check_inputs (line 205) | def check_inputs(self, num_inference_steps, timesteps, latents, batch_...
method __call__ (line 230) | def __call__(
FILE: 4DoF/diffusers/pipelines/controlnet/multicontrolnet.py
class MultiControlNetModel (line 15) | class MultiControlNetModel(ModelMixin):
method __init__ (line 28) | def __init__(self, controlnets: Union[List[ControlNetModel], Tuple[Con...
method forward (line 32) | def forward(
method save_pretrained (line 73) | def save_pretrained(
method from_pretrained (line 116) | def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os...
FILE: 4DoF/diffusers/pipelines/controlnet/pipeline_controlnet.py
class StableDiffusionControlNetPipeline (line 93) | class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInvers...
method __init__ (line 129) | def __init__(
method enable_vae_slicing (line 180) | def enable_vae_slicing(self):
method disable_vae_slicing (line 190) | def disable_vae_slicing(self):
method enable_vae_tiling (line 198) | def enable_vae_tiling(self):
method disable_vae_tiling (line 208) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 215) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 236) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 266) | def _execution_device(self):
method _encode_prompt (line 284) | def _encode_prompt(
method run_safety_checker (line 439) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 454) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 468) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 485) | def check_inputs(
method check_image (line 619) | def check_image(self, image, prompt, prompt_embeds):
method prepare_image (line 656) | def prepare_image(
method prepare_latents (line 687) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 706) | def __call__(
FILE: 4DoF/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py
function prepare_image (line 95) | def prepare_image(image):
class StableDiffusionControlNetImg2ImgPipeline (line 119) | class StableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline, Textua...
method __init__ (line 155) | def __init__(
method enable_vae_slicing (line 206) | def enable_vae_slicing(self):
method disable_vae_slicing (line 216) | def disable_vae_slicing(self):
method enable_vae_tiling (line 224) | def enable_vae_tiling(self):
method disable_vae_tiling (line 234) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 241) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 262) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 292) | def _execution_device(self):
method _encode_prompt (line 310) | def _encode_prompt(
method run_safety_checker (line 465) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 480) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 494) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 511) | def check_inputs(
method check_image (line 646) | def check_image(self, image, prompt, prompt_embeds):
method prepare_control_image (line 684) | def prepare_control_image(
method get_timesteps (line 715) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 725) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method __call__ (line 784) | def __call__(
FILE: 4DoF/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py
function prepare_mask_and_masked_image (line 108) | def prepare_mask_and_masked_image(image, mask, height, width, return_ima...
class StableDiffusionControlNetInpaintPipeline (line 225) | class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline, Textua...
method __init__ (line 272) | def __init__(
method enable_vae_slicing (line 323) | def enable_vae_slicing(self):
method disable_vae_slicing (line 333) | def disable_vae_slicing(self):
method enable_vae_tiling (line 341) | def enable_vae_tiling(self):
method disable_vae_tiling (line 351) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 358) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 379) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 409) | def _execution_device(self):
method _encode_prompt (line 427) | def _encode_prompt(
method run_safety_checker (line 582) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 597) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 611) | def prepare_extra_step_kwargs(self, generator, eta):
method get_timesteps (line 629) | def get_timesteps(self, num_inference_steps, strength, device):
method check_inputs (line 638) | def check_inputs(
method check_image (line 778) | def check_image(self, image, prompt, prompt_embeds):
method prepare_control_image (line 816) | def prepare_control_image(
method prepare_latents (line 847) | def prepare_latents(
method _default_height_width (line 900) | def _default_height_width(self, height, width, image):
method prepare_mask_latents (line 926) | def prepare_mask_latents(
method _encode_vae_image (line 968) | def _encode_vae_image(self, image: torch.Tensor, generator: torch.Gene...
method __call__ (line 984) | def __call__(
FILE: 4DoF/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py
class FlaxStableDiffusionControlNetPipeline (line 119) | class FlaxStableDiffusionControlNetPipeline(FlaxDiffusionPipeline):
method __init__ (line 150) | def __init__(
method prepare_text_inputs (line 189) | def prepare_text_inputs(self, prompt: Union[str, List[str]]):
method prepare_image_inputs (line 203) | def prepare_image_inputs(self, image: Union[Image.Image, List[Image.Im...
method _get_has_nsfw_concepts (line 214) | def _get_has_nsfw_concepts(self, features, params):
method _run_safety_checker (line 218) | def _run_safety_checker(self, images, safety_model_params, jit=False):
method _generate (line 248) | def _generate(
method __call__ (line 358) | def __call__(
function _p_generate (line 493) | def _p_generate(
function _p_get_has_nsfw_concepts (line 519) | def _p_get_has_nsfw_concepts(pipe, features, params):
function unshard (line 523) | def unshard(x: jnp.ndarray):
function preprocess (line 530) | def preprocess(image, dtype):
FILE: 4DoF/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py
class DanceDiffusionPipeline (line 27) | class DanceDiffusionPipeline(DiffusionPipeline):
method __init__ (line 39) | def __init__(self, unet, scheduler):
method __call__ (line 44) | def __call__(
FILE: 4DoF/diffusers/pipelines/ddim/pipeline_ddim.py
class DDIMPipeline (line 24) | class DDIMPipeline(DiffusionPipeline):
method __init__ (line 36) | def __init__(self, unet, scheduler):
method __call__ (line 45) | def __call__(
FILE: 4DoF/diffusers/pipelines/ddpm/pipeline_ddpm.py
class DDPMPipeline (line 24) | class DDPMPipeline(DiffusionPipeline):
method __init__ (line 36) | def __init__(self, unet, scheduler):
method __call__ (line 41) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/__init__.py
class IFPipelineOutput (line 21) | class IFPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if.py
class IFPipeline (line 89) | class IFPipeline(DiffusionPipeline, LoraLoaderMixin):
method __init__ (line 107) | def __init__(
method enable_sequential_cpu_offload (line 147) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 171) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 212) | def remove_all_hooks(self):
method _execution_device (line 228) | def _execution_device(self):
method encode_prompt (line 246) | def encode_prompt(
method run_safety_checker (line 397) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 414) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 431) | def check_inputs(
method prepare_intermediate_images (line 473) | def prepare_intermediate_images(self, batch_size, num_channels, height...
method _text_preprocessing (line 487) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 511) | def _clean_caption(self, caption):
method __call__ (line 627) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py
function resize (line 41) | def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image:
class IFImg2ImgPipeline (line 113) | class IFImg2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
method __init__ (line 131) | def __init__(
method enable_sequential_cpu_offload (line 172) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 197) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 239) | def remove_all_hooks(self):
method _execution_device (line 255) | def _execution_device(self):
method encode_prompt (line 274) | def encode_prompt(
method run_safety_checker (line 426) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 443) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 460) | def check_inputs(
method _text_preprocessing (line 534) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 559) | def _clean_caption(self, caption):
method preprocess_image (line 673) | def preprocess_image(self, image: PIL.Image.Image) -> torch.Tensor:
method get_timesteps (line 709) | def get_timesteps(self, num_inference_steps, strength):
method prepare_intermediate_images (line 718) | def prepare_intermediate_images(
method __call__ (line 742) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py
function resize (line 44) | def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image:
class IFImg2ImgSuperResolutionPipeline (line 116) | class IFImg2ImgSuperResolutionPipeline(DiffusionPipeline, LoraLoaderMixin):
method __init__ (line 135) | def __init__(
method enable_sequential_cpu_offload (line 183) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 208) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 250) | def remove_all_hooks(self):
method _text_preprocessing (line 265) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 290) | def _clean_caption(self, caption):
method _execution_device (line 406) | def _execution_device(self):
method encode_prompt (line 425) | def encode_prompt(
method run_safety_checker (line 577) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 594) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 611) | def check_inputs(
method preprocess_original_image (line 721) | def preprocess_original_image(self, image: PIL.Image.Image) -> torch.T...
method preprocess_image (line 758) | def preprocess_image(self, image: PIL.Image.Image, num_images_per_prom...
method get_timesteps (line 790) | def get_timesteps(self, num_inference_steps, strength):
method prepare_intermediate_images (line 800) | def prepare_intermediate_images(
method __call__ (line 824) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py
function resize (line 42) | def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image:
class IFInpaintingPipeline (line 116) | class IFInpaintingPipeline(DiffusionPipeline, LoraLoaderMixin):
method __init__ (line 134) | def __init__(
method enable_sequential_cpu_offload (line 175) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 200) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 242) | def remove_all_hooks(self):
method _execution_device (line 258) | def _execution_device(self):
method encode_prompt (line 277) | def encode_prompt(
method run_safety_checker (line 429) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 446) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 463) | def check_inputs(
method _text_preprocessing (line 573) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 598) | def _clean_caption(self, caption):
method preprocess_image (line 713) | def preprocess_image(self, image: PIL.Image.Image) -> torch.Tensor:
method preprocess_mask_image (line 749) | def preprocess_mask_image(self, mask_image) -> torch.Tensor:
method get_timesteps (line 799) | def get_timesteps(self, num_inference_steps, strength):
method prepare_intermediate_images (line 808) | def prepare_intermediate_images(
method __call__ (line 834) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py
function resize (line 44) | def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image:
class IFInpaintingSuperResolutionPipeline (line 118) | class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, LoraLoaderM...
method __init__ (line 137) | def __init__(
method enable_sequential_cpu_offload (line 185) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 210) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 252) | def remove_all_hooks(self):
method _text_preprocessing (line 267) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 292) | def _clean_caption(self, caption):
method _execution_device (line 408) | def _execution_device(self):
method encode_prompt (line 427) | def encode_prompt(
method run_safety_checker (line 579) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 596) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 613) | def check_inputs(
method preprocess_original_image (line 757) | def preprocess_original_image(self, image: PIL.Image.Image) -> torch.T...
method preprocess_image (line 794) | def preprocess_image(self, image: PIL.Image.Image, num_images_per_prom...
method preprocess_mask_image (line 826) | def preprocess_mask_image(self, mask_image) -> torch.Tensor:
method get_timesteps (line 876) | def get_timesteps(self, num_inference_steps, strength):
method prepare_intermediate_images (line 886) | def prepare_intermediate_images(
method __call__ (line 912) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py
class IFSuperResolutionPipeline (line 74) | class IFSuperResolutionPipeline(DiffusionPipeline, LoraLoaderMixin):
method __init__ (line 93) | def __init__(
method enable_sequential_cpu_offload (line 141) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 166) | def enable_model_cpu_offload(self, gpu_id=0):
method remove_all_hooks (line 208) | def remove_all_hooks(self):
method _text_preprocessing (line 223) | def _text_preprocessing(self, text, clean_caption=False):
method _clean_caption (line 248) | def _clean_caption(self, caption):
method _execution_device (line 364) | def _execution_device(self):
method encode_prompt (line 383) | def encode_prompt(
method run_safety_checker (line 535) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 552) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 569) | def check_inputs(
method prepare_intermediate_images (line 649) | def prepare_intermediate_images(self, batch_size, num_channels, height...
method preprocess_image (line 663) | def preprocess_image(self, image, num_images_per_prompt, device):
method __call__ (line 696) | def __call__(
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/safety_checker.py
class IFSafetyChecker (line 12) | class IFSafetyChecker(PreTrainedModel):
method __init__ (line 17) | def __init__(self, config: CLIPConfig):
method forward (line 26) | def forward(self, clip_input, images, p_threshold=0.5, w_threshold=0.5):
FILE: 4DoF/diffusers/pipelines/deepfloyd_if/watermark.py
class IFWatermarker (line 12) | class IFWatermarker(ModelMixin, ConfigMixin):
method __init__ (line 13) | def __init__(self):
method apply_watermark (line 19) | def apply_watermark(self, images: List[PIL.Image.Image], sample_size=N...
FILE: 4DoF/diffusers/pipelines/dit/pipeline_dit.py
class DiTPipeline (line 31) | class DiTPipeline(DiffusionPipeline):
method __init__ (line 45) | def __init__(
method get_label_ids (line 63) | def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
method __call__ (line 87) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky.py
function get_new_h_w (line 69) | def get_new_h_w(h, w, scale_factor=8):
class KandinskyPipeline (line 79) | class KandinskyPipeline(DiffusionPipeline):
method __init__ (line 99) | def __init__(
method prepare_latents (line 119) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method _encode_prompt (line 130) | def _encode_prompt(
method enable_sequential_cpu_offload (line 232) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 254) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 284) | def _execution_device(self):
method __call__ (line 303) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py
function get_new_h_w (line 82) | def get_new_h_w(h, w, scale_factor=8):
function prepare_image (line 92) | def prepare_image(pil_image, w=512, h=512):
class KandinskyImg2ImgPipeline (line 101) | class KandinskyImg2ImgPipeline(DiffusionPipeline):
method __init__ (line 121) | def __init__(
method get_timesteps (line 140) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 149) | def prepare_latents(self, latents, latent_timestep, shape, dtype, devi...
method _encode_prompt (line 165) | def _encode_prompt(
method enable_sequential_cpu_offload (line 267) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 289) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 319) | def _execution_device(self):
method add_noise (line 337) | def add_noise(
method __call__ (line 365) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py
function get_new_h_w (line 89) | def get_new_h_w(h, w, scale_factor=8):
function prepare_mask (line 99) | def prepare_mask(masks):
function prepare_mask_and_masked_image (line 123) | def prepare_mask_and_masked_image(image, mask, height, width):
class KandinskyInpaintPipeline (line 239) | class KandinskyInpaintPipeline(DiffusionPipeline):
method __init__ (line 259) | def __init__(
method prepare_latents (line 279) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method _encode_prompt (line 290) | def _encode_prompt(
method enable_sequential_cpu_offload (line 392) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 414) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 444) | def _execution_device(self):
method __call__ (line 463) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py
class KandinskyPriorPipelineOutput (line 114) | class KandinskyPriorPipelineOutput(BaseOutput):
class KandinskyPriorPipeline (line 129) | class KandinskyPriorPipeline(DiffusionPipeline):
method __init__ (line 150) | def __init__(
method interpolate (line 172) | def interpolate(
method prepare_latents (line 278) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method get_zero_embed (line 289) | def get_zero_embed(self, batch_size=1, device=None):
method enable_sequential_cpu_offload (line 298) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 320) | def _execution_device(self):
method _encode_prompt (line 337) | def _encode_prompt(
method __call__ (line 436) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky/text_encoder.py
class MCLIPConfig (line 5) | class MCLIPConfig(XLMRobertaConfig):
method __init__ (line 8) | def __init__(self, transformerDimSize=1024, imageDimSize=768, **kwargs):
class MultilingualCLIP (line 14) | class MultilingualCLIP(PreTrainedModel):
method __init__ (line 17) | def __init__(self, config, *args, **kwargs):
method forward (line 24) | def forward(self, input_ids, attention_mask):
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py
function downscale_height_and_width (line 60) | def downscale_height_and_width(height, width, scale_factor=8):
class KandinskyV22Pipeline (line 70) | class KandinskyV22Pipeline(DiffusionPipeline):
method __init__ (line 86) | def __init__(
method prepare_latents (line 102) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method enable_sequential_cpu_offload (line 113) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 134) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 161) | def _execution_device(self):
method __call__ (line 180) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py
function downscale_height_and_width (line 100) | def downscale_height_and_width(height, width, scale_factor=8):
class KandinskyV22ControlnetPipeline (line 110) | class KandinskyV22ControlnetPipeline(DiffusionPipeline):
method __init__ (line 126) | def __init__(
method prepare_latents (line 142) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method enable_sequential_cpu_offload (line 154) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 176) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 203) | def _execution_device(self):
method __call__ (line 222) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py
function downscale_height_and_width (line 104) | def downscale_height_and_width(height, width, scale_factor=8):
function prepare_image (line 115) | def prepare_image(pil_image, w=512, h=512):
class KandinskyV22ControlnetImg2ImgPipeline (line 124) | class KandinskyV22ControlnetImg2ImgPipeline(DiffusionPipeline):
method __init__ (line 140) | def __init__(
method get_timesteps (line 156) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 166) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method enable_sequential_cpu_offload (line 209) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 231) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 258) | def _execution_device(self):
method __call__ (line 277) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py
function downscale_height_and_width (line 78) | def downscale_height_and_width(height, width, scale_factor=8):
function prepare_image (line 89) | def prepare_image(pil_image, w=512, h=512):
class KandinskyV22Img2ImgPipeline (line 98) | class KandinskyV22Img2ImgPipeline(DiffusionPipeline):
method __init__ (line 114) | def __init__(
method get_timesteps (line 130) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 139) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method enable_sequential_cpu_offload (line 182) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 204) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 231) | def _execution_device(self):
method __call__ (line 250) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py
function downscale_height_and_width (line 85) | def downscale_height_and_width(height, width, scale_factor=8):
function prepare_mask (line 96) | def prepare_mask(masks):
function prepare_mask_and_masked_image (line 121) | def prepare_mask_and_masked_image(image, mask, height, width):
class KandinskyV22InpaintPipeline (line 237) | class KandinskyV22InpaintPipeline(DiffusionPipeline):
method __init__ (line 253) | def __init__(
method prepare_latents (line 269) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method enable_sequential_cpu_offload (line 281) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 303) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 330) | def _execution_device(self):
method __call__ (line 349) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py
class KandinskyV22PriorPipeline (line 85) | class KandinskyV22PriorPipeline(DiffusionPipeline):
method __init__ (line 108) | def __init__(
method interpolate (line 130) | def interpolate(
method prepare_latents (line 236) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method get_zero_embed (line 248) | def get_zero_embed(self, batch_size=1, device=None):
method enable_sequential_cpu_offload (line 258) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 281) | def _execution_device(self):
method _encode_prompt (line 299) | def _encode_prompt(
method __call__ (line 398) | def __call__(
FILE: 4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py
class KandinskyV22PriorEmb2EmbPipeline (line 103) | class KandinskyV22PriorEmb2EmbPipeline(DiffusionPipeline):
method __init__ (line 124) | def __init__(
method get_timesteps (line 144) | def get_timesteps(self, num_inference_steps, strength, device):
method interpolate (line 155) | def interpolate(
method _encode_image (line 243) | def _encode_image(
method prepare_latents (line 260) | def prepare_latents(self, emb, timestep, batch_size, num_images_per_pr...
method get_zero_embed (line 287) | def get_zero_embed(self, batch_size=1, device=None):
method enable_sequential_cpu_offload (line 297) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 320) | def _execution_device(self):
method _encode_prompt (line 338) | def _encode_prompt(
method __call__ (line 437) | def __call__(
FILE: 4DoF/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
class LDMTextToImagePipeline (line 32) | class LDMTextToImagePipeline(DiffusionPipeline):
method __init__ (line 51) | def __init__(
method __call__ (line 64) | def __call__(
class LDMBertConfig (line 220) | class LDMBertConfig(PretrainedConfig):
method __init__ (line 225) | def __init__(
function _expand_mask (line 267) | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Option...
class LDMBertAttention (line 282) | class LDMBertAttention(nn.Module):
method __init__ (line 285) | def __init__(
method _shape (line 309) | def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
method forward (line 312) | def forward(
class LDMBertEncoderLayer (line 426) | class LDMBertEncoderLayer(nn.Module):
method __init__ (line 427) | def __init__(self, config: LDMBertConfig):
method forward (line 444) | def forward(
class LDMBertPreTrainedModel (line 496) | class LDMBertPreTrainedModel(PreTrainedModel):
method _init_weights (line 502) | def _init_weights(self, module):
method _set_gradient_checkpointing (line 513) | def _set_gradient_checkpointing(self, module, value=False):
method dummy_inputs (line 518) | def dummy_inputs(self):
class LDMBertEncoder (line 528) | class LDMBertEncoder(LDMBertPreTrainedModel):
method __init__ (line 538) | def __init__(self, config: LDMBertConfig):
method get_input_embeddings (line 556) | def get_input_embeddings(self):
method set_input_embeddings (line 559) | def set_input_embeddings(self, value):
method forward (line 562) | def forward(
class LDMBertModel (line 695) | class LDMBertModel(LDMBertPreTrainedModel):
method __init__ (line 698) | def __init__(self, config: LDMBertConfig):
method forward (line 703) | def forward(
FILE: 4DoF/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
function preprocess (line 22) | def preprocess(image):
class LDMSuperResolutionPipeline (line 32) | class LDMSuperResolutionPipeline(DiffusionPipeline):
method __init__ (line 49) | def __init__(
method __call__ (line 66) | def __call__(
FILE: 4DoF/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py
class LDMPipeline (line 26) | class LDMPipeline(DiffusionPipeline):
method __init__ (line 39) | def __init__(self, vqvae: VQModel, unet: UNet2DModel, scheduler: DDIMS...
method __call__ (line 44) | def __call__(
FILE: 4DoF/diffusers/pipelines/onnx_utils.py
class OnnxRuntimeModel (line 51) | class OnnxRuntimeModel:
method __init__ (line 52) | def __init__(self, model=None, **kwargs):
method __call__ (line 58) | def __call__(self, **kwargs):
method load_model (line 63) | def load_model(path: Union[str, Path], provider=None, sess_options=None):
method _save_pretrained (line 79) | def _save_pretrained(self, save_directory: Union[str, Path], file_name...
method save_pretrained (line 110) | def save_pretrained(
method _from_pretrained (line 133) | def _from_pretrained(
method from_pretrained (line 193) | def from_pretrained(
FILE: 4DoF/diffusers/pipelines/paint_by_example/image_encoder.py
class PaintByExampleImageEncoder (line 25) | class PaintByExampleImageEncoder(CLIPPreTrainedModel):
method __init__ (line 26) | def __init__(self, config, proj_size=768):
method forward (line 38) | def forward(self, pixel_values, return_uncond_vector=False):
class PaintByExampleMapper (line 50) | class PaintByExampleMapper(nn.Module):
method __init__ (line 51) | def __init__(self, config):
method forward (line 63) | def forward(self, hidden_states):
FILE: 4DoF/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
function prepare_mask_and_masked_image (line 39) | def prepare_mask_and_masked_image(image, mask):
class PaintByExamplePipeline (line 139) | class PaintByExamplePipeline(DiffusionPipeline):
method __init__ (line 168) | def __init__(
method enable_sequential_cpu_offload (line 192) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 213) | def _execution_device(self):
method run_safety_checker (line 231) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 246) | def prepare_extra_step_kwargs(self, generator, eta):
method decode_latents (line 264) | def decode_latents(self, latents):
method check_inputs (line 278) | def check_inputs(self, image, height, width, callback_steps):
method prepare_latents (line 301) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method prepare_mask_latents (line 319) | def prepare_mask_latents(
method _encode_vae_image (line 361) | def _encode_vae_image(self, image: torch.Tensor, generator: torch.Gene...
method _encode_image (line 375) | def _encode_image(self, image, device, num_images_per_prompt, do_class...
method __call__ (line 401) | def __call__(
FILE: 4DoF/diffusers/pipelines/pipeline_flax_utils.py
function import_flax_or_no_model (line 66) | def import_flax_or_no_model(module, class_name):
class FlaxImagePipelineOutput (line 80) | class FlaxImagePipelineOutput(BaseOutput):
class FlaxDiffusionPipeline (line 93) | class FlaxDiffusionPipeline(ConfigMixin):
method register_modules (line 110) | def register_modules(self, **kwargs):
method save_pretrained (line 143) | def save_pretrained(self, save_directory: Union[str, os.PathLike], par...
method from_pretrained (line 194) | def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union...
method _get_signature_keys (line 496) | def _get_signature_keys(obj):
method components (line 504) | def components(self) -> Dict[str, Any]:
method numpy_to_pil (line 541) | def numpy_to_pil(images):
method progress_bar (line 557) | def progress_bar(self, iterable):
method set_progress_bar_config (line 567) | def set_progress_bar_config(self, **kwargs):
FILE: 4DoF/diffusers/pipelines/pipeline_utils.py
class ImagePipelineOutput (line 112) | class ImagePipelineOutput(BaseOutput):
class AudioPipelineOutput (line 126) | class AudioPipelineOutput(BaseOutput):
function is_safetensors_compatible (line 138) | def is_safetensors_compatible(filenames, variant=None, passed_components...
function variant_compatible_siblings (line 187) | def variant_compatible_siblings(filenames, variant=None) -> Union[List[o...
function warn_deprecated_model_variant (line 254) | def warn_deprecated_model_variant(pretrained_model_name_or_path, use_aut...
function maybe_raise_or_warn (line 276) | def maybe_raise_or_warn(
function get_class_obj_and_candidates (line 308) | def get_class_obj_and_candidates(library_name, class_name, importable_cl...
function _get_pipeline_class (line 325) | def _get_pipeline_class(class_obj, config, custom_pipeline=None, cache_d...
function load_sub_model (line 346) | def load_sub_model(
class DiffusionPipeline (line 458) | class DiffusionPipeline(ConfigMixin):
method register_modules (line 478) | def register_modules(self, **kwargs):
method __setattr__ (line 521) | def __setattr__(self, name: str, value: Any):
method save_pretrained (line 536) | def save_pretrained(
method to (line 623) | def to(
method device (line 699) | def device(self) -> torch.device:
method from_pretrained (line 714) | def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union...
method download (line 1079) | def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.Pa...
method _get_signature_keys (line 1340) | def _get_signature_keys(obj):
method components (line 1348) | def components(self) -> Dict[str, Any]:
method numpy_to_pil (line 1384) | def numpy_to_pil(images):
method progress_bar (line 1390) | def progress_bar(self, iterable=None, total=None):
method set_progress_bar_config (line 1405) | def set_progress_bar_config(self, **kwargs):
method enable_xformers_memory_efficient_attention (line 1408) | def enable_xformers_memory_efficient_attention(self, attention_op: Opt...
method disable_xformers_memory_efficient_attention (line 1444) | def disable_xformers_memory_efficient_attention(self):
method set_use_memory_efficient_attention_xformers (line 1450) | def set_use_memory_efficient_attention_xformers(
method enable_attention_slicing (line 1470) | def enable_attention_slicing(self, slice_size: Optional[Union[str, int...
method disable_attention_slicing (line 1486) | def disable_attention_slicing(self):
method set_attention_slice (line 1494) | def set_attention_slice(self, slice_size: Optional[int]):
FILE: 4DoF/diffusers/pipelines/pndm/pipeline_pndm.py
class PNDMPipeline (line 26) | class PNDMPipeline(DiffusionPipeline):
method __init__ (line 40) | def __init__(self, unet: UNet2DModel, scheduler: PNDMScheduler):
method __call__ (line 48) | def __call__(
FILE: 4DoF/diffusers/pipelines/repaint/pipeline_repaint.py
function _preprocess_image (line 33) | def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]):
function _preprocess_mask (line 59) | def _preprocess_mask(mask: Union[List, PIL.Image.Image, torch.Tensor]):
class RePaintPipeline (line 79) | class RePaintPipeline(DiffusionPipeline):
method __init__ (line 83) | def __init__(self, unet, scheduler):
method __call__ (line 88) | def __call__(
FILE: 4DoF/diffusers/pipelines/score_sde_ve/pipeline_score_sde_ve.py
class ScoreSdeVePipeline (line 25) | class ScoreSdeVePipeline(DiffusionPipeline):
method __init__ (line 36) | def __init__(self, unet: UNet2DModel, scheduler: ScoreSdeVeScheduler):
method __call__ (line 41) | def __call__(
FILE: 4DoF/diffusers/pipelines/semantic_stable_diffusion/__init__.py
class SemanticStableDiffusionPipelineOutput (line 13) | class SemanticStableDiffusionPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py
class SemanticStableDiffusionPipeline (line 63) | class SemanticStableDiffusionPipeline(DiffusionPipeline):
method __init__ (line 95) | def __init__(
method run_safety_checker (line 138) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 153) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 167) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 185) | def check_inputs(
method prepare_latents (line 233) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 251) | def __call__(
FILE: 4DoF/diffusers/pipelines/shap_e/camera.py
class DifferentiableProjectiveCamera (line 23) | class DifferentiableProjectiveCamera:
method __post_init__ (line 38) | def __post_init__(self):
method resolution (line 43) | def resolution(self):
method fov (line 46) | def fov(self):
method get_image_coords (line 49) | def get_image_coords(self) -> torch.Tensor:
method camera_rays (line 64) | def camera_rays(self):
method get_camera_rays (line 76) | def get_camera_rays(self, coords: torch.Tensor) -> torch.Tensor:
method resize_image (line 105) | def resize_image(self, width: int, height: int) -> "DifferentiableProj...
function create_pan_cameras (line 122) | def create_pan_cameras(size: int) -> DifferentiableProjectiveCamera:
FILE: 4DoF/diffusers/pipelines/shap_e/pipeline_shap_e.py
class ShapEPipelineOutput (line 69) | class ShapEPipelineOutput(BaseOutput):
class ShapEPipeline (line 81) | class ShapEPipeline(DiffusionPipeline):
method __init__ (line 103) | def __init__(
method prepare_latents (line 122) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method enable_sequential_cpu_offload (line 133) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 151) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 180) | def _execution_device(self):
method _encode_prompt (line 197) | def _encode_prompt(
method __call__ (line 248) | def __call__(
FILE: 4DoF/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py
class ShapEPipelineOutput (line 69) | class ShapEPipelineOutput(BaseOutput):
class ShapEImg2ImgPipeline (line 81) | class ShapEImg2ImgPipeline(DiffusionPipeline):
method __init__ (line 103) | def __init__(
method prepare_latents (line 122) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method enable_sequential_cpu_offload (line 133) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 152) | def _execution_device(self):
method _encode_image (line 169) | def _encode_image(
method __call__ (line 201) | def __call__(
FILE: 4DoF/diffusers/pipelines/shap_e/renderer.py
function sample_pmf (line 30) | def sample_pmf(pmf: torch.Tensor, n_samples: int) -> torch.Tensor:
function posenc_nerf (line 53) | def posenc_nerf(x: torch.Tensor, min_deg: int = 0, max_deg: int = 15) ->...
function encode_position (line 70) | def encode_position(position):
function encode_direction (line 74) | def encode_direction(position, direction=None):
function _sanitize_name (line 81) | def _sanitize_name(x: str) -> str:
function integrate_samples (line 85) | def integrate_samples(volume_range, ts, density, channels):
class VoidNeRFModel (line 119) | class VoidNeRFModel(nn.Module):
method __init__ (line 124) | def __init__(self, background, channel_scale=255.0):
method forward (line 130) | def forward(self, position):
class VolumeRange (line 142) | class VolumeRange:
method __post_init__ (line 147) | def __post_init__(self):
method partition (line 150) | def partition(self, ts):
class BoundingBoxVolume (line 174) | class BoundingBoxVolume(nn.Module):
method __init__ (line 179) | def __init__(
method intersect (line 205) | def intersect(
class StratifiedRaySampler (line 258) | class StratifiedRaySampler(nn.Module):
method __init__ (line 263) | def __init__(self, depth_mode: str = "linear"):
method sample (line 271) | def sample(
class ImportanceRaySampler (line 310) | class ImportanceRaySampler(nn.Module):
method __init__ (line 315) | def __init__(
method sample (line 338) | def sample(self, t0: torch.Tensor, t1: torch.Tensor, n_samples: int) -...
class MLPNeRFModelOutput (line 372) | class MLPNeRFModelOutput(BaseOutput):
class MLPNeRSTFModel (line 379) | class MLPNeRSTFModel(ModelMixin, ConfigMixin):
method __init__ (line 381) | def __init__(
method map_indices_to_keys (line 418) | def map_indices_to_keys(self, output):
method forward (line 432) | def forward(self, *, position, direction, ts, nerf_level="coarse"):
class ChannelsProj (line 471) | class ChannelsProj(nn.Module):
method __init__ (line 472) | def __init__(
method forward (line 486) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class ShapEParamsProjModel (line 497) | class ShapEParamsProjModel(ModelMixin, ConfigMixin):
method __init__ (line 505) | def __init__(
method forward (line 535) | def forward(self, x: torch.Tensor):
class ShapERenderer (line 547) | class ShapERenderer(ModelMixin, ConfigMixin):
method __init__ (line 549) | def __init__(
method render_rays (line 588) | def render_rays(self, rays, sampler, n_samples, prev_model_out=None, r...
method decode (line 667) | def decode(
FILE: 4DoF/diffusers/pipelines/spectrogram_diffusion/continous_encoder.py
class SpectrogramContEncoder (line 29) | class SpectrogramContEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin):
method __init__ (line 31) | def __init__(
method forward (line 71) | def forward(self, encoder_inputs, encoder_inputs_mask):
FILE: 4DoF/diffusers/pipelines/spectrogram_diffusion/midi_utils.py
class NoteRepresentationConfig (line 84) | class NoteRepresentationConfig:
class NoteEventData (line 92) | class NoteEventData:
class NoteEncodingState (line 101) | class NoteEncodingState:
class EventRange (line 109) | class EventRange:
class Event (line 116) | class Event:
class Tokenizer (line 121) | class Tokenizer:
method __init__ (line 122) | def __init__(self, regular_ids: int):
method encode (line 127) | def encode(self, token_ids):
class Codec (line 145) | class Codec:
method __init__ (line 156) | def __init__(self, max_shift_steps: int, steps_per_second: float, even...
method num_classes (line 172) | def num_classes(self) -> int:
method is_shift_event_index (line 178) | def is_shift_event_index(self, index: int) -> bool:
method max_shift_steps (line 182) | def max_shift_steps(self) -> int:
method encode_event (line 185) | def encode_event(self, event: Event) -> int:
method event_type_range (line 200) | def event_type_range(self, event_type: str) -> Tuple[int, int]:
method decode_event_index (line 210) | def decode_event_index(self, index: int) -> Event:
class ProgramGranularity (line 222) | class ProgramGranularity:
function drop_programs (line 228) | def drop_programs(tokens, codec: Codec):
function programs_to_midi_classes (line 234) | def programs_to_midi_classes(tokens, codec):
function frame (line 254) | def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, ...
function program_to_slakh_program (line 272) | def program_to_slakh_program(program):
function audio_to_frames (line 279) | def audio_to_frames(
function note_sequence_to_onsets_and_offsets_and_programs (line 302) | def note_sequence_to_onsets_and_offsets_and_programs(
function num_velocity_bins_from_codec (line 332) | def num_velocity_bins_from_codec(codec: Codec):
function segment (line 339) | def segment(a, n):
function velocity_to_bin (line 343) | def velocity_to_bin(velocity, num_velocity_bins):
function note_event_data_to_events (line 350) | def note_event_data_to_events(
function note_encoding_state_to_events (line 382) | def note_encoding_state_to_events(state: NoteEncodingState) -> Sequence[...
function encode_and_index_events (line 392) | def encode_and_index_events(
function extract_sequence_with_indices (line 498) | def extract_sequence_with_indices(features, state_events_end_token=None,...
function map_midi_programs (line 524) | def map_midi_programs(
function run_length_encode_shifts_fn (line 534) | def run_length_encode_shifts_fn(
function note_representation_processor_chain (line 604) | def note_representation_processor_chain(features, codec: Codec, note_rep...
class MidiProcessor (line 619) | class MidiProcessor:
method __init__ (line 620) | def __init__(self):
method __call__ (line 635) | def __call__(self, midi: Union[bytes, os.PathLike, str]):
FILE: 4DoF/diffusers/pipelines/spectrogram_diffusion/notes_encoder.py
class SpectrogramNotesEncoder (line 25) | class SpectrogramNotesEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin):
method __init__ (line 27) | def __init__(
method forward (line 69) | def forward(self, encoder_input_tokens, encoder_inputs_mask):
FILE: 4DoF/diffusers/pipelines/spectrogram_diffusion/pipeline_spectrogram_diffusion.py
class SpectrogramDiffusionPipeline (line 40) | class SpectrogramDiffusionPipeline(DiffusionPipeline):
method __init__ (line 43) | def __init__(
method scale_features (line 66) | def scale_features(self, features, output_range=(-1.0, 1.0), clip=False):
method scale_to_features (line 76) | def scale_to_features(self, outputs, input_range=(-1.0, 1.0), clip=Fal...
method encode (line 85) | def encode(self, input_tokens, continuous_inputs, continuous_mask):
method decode (line 97) | def decode(self, encodings_and_masks, input_tokens, noise_time):
method __call__ (line 113) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/__init__.py
class StableDiffusionPipelineOutput (line 22) | class StableDiffusionPipelineOutput(BaseOutput):
class FlaxStableDiffusionPipelineOutput (line 117) | class FlaxStableDiffusionPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
function shave_segments (line 68) | def shave_segments(path, n_shave_prefix_segments=1):
function renew_resnet_paths (line 78) | def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
function renew_vae_resnet_paths (line 100) | def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
function renew_attention_paths (line 116) | def renew_attention_paths(old_list, n_shave_prefix_segments=0):
function renew_vae_attention_paths (line 137) | def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
function assign_to_checkpoint (line 167) | def assign_to_checkpoint(
function conv_attn_to_linear (line 222) | def conv_attn_to_linear(checkpoint):
function create_unet_diffusers_config (line 234) | def create_unet_diffusers_config(original_config, image_size: int, contr...
function create_vae_diffusers_config (line 334) | def create_vae_diffusers_config(original_config, image_size: int):
function create_diffusers_schedular (line 358) | def create_diffusers_schedular(original_config):
function create_ldm_bert_config (line 368) | def create_ldm_bert_config(original_config):
function convert_ldm_unet_checkpoint (line 378) | def convert_ldm_unet_checkpoint(
function convert_ldm_vae_checkpoint (line 620) | def convert_ldm_vae_checkpoint(checkpoint, config):
function convert_ldm_bert_checkpoint (line 727) | def convert_ldm_bert_checkpoint(checkpoint, config):
function convert_ldm_clip_checkpoint (line 777) | def convert_ldm_clip_checkpoint(checkpoint, local_files_only=False, text...
function convert_paint_by_example_checkpoint (line 827) | def convert_paint_by_example_checkpoint(checkpoint):
function convert_open_clip_checkpoint (line 894) | def convert_open_clip_checkpoint(
function stable_unclip_image_encoder (line 959) | def stable_unclip_image_encoder(original_config):
function stable_unclip_image_noising_components (line 992) | def stable_unclip_image_noising_components(
function convert_controlnet_checkpoint (line 1037) | def convert_controlnet_checkpoint(
function download_from_original_stable_diffusion_ckpt (line 1081) | def download_from_original_stable_diffusion_ckpt(
function download_controlnet_from_original_ckpt (line 1563) | def download_controlnet_from_original_ckpt(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py
function preprocess (line 42) | def preprocess(image):
function posterior_sample (line 68) | def posterior_sample(scheduler, latents, timestep, clean_latents, genera...
function compute_noise (line 95) | def compute_noise(scheduler, prev_latents, latents, timestep, noise_pred...
class CycleDiffusionPipeline (line 129) | class CycleDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMi...
method __init__ (line 158) | def __init__(
method enable_sequential_cpu_offload (line 235) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 261) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 291) | def _execution_device(self):
method _encode_prompt (line 309) | def _encode_prompt(
method check_inputs (line 464) | def check_inputs(
method prepare_extra_step_kwargs (line 505) | def prepare_extra_step_kwargs(self, generator, eta):
method run_safety_checker (line 523) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 538) | def decode_latents(self, latents):
method get_timesteps (line 552) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 561) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method __call__ (line 616) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
class FlaxStableDiffusionPipeline (line 81) | class FlaxStableDiffusionPipeline(FlaxDiffusionPipeline):
method __init__ (line 110) | def __init__(
method prepare_inputs (line 168) | def prepare_inputs(self, prompt: Union[str, List[str]]):
method _get_has_nsfw_concepts (line 181) | def _get_has_nsfw_concepts(self, features, params):
method _run_safety_checker (line 185) | def _run_safety_checker(self, images, safety_model_params, jit=False):
method _generate (line 215) | def _generate(
method __call__ (line 312) | def __call__(
function _p_generate (line 436) | def _p_generate(
function _p_get_has_nsfw_concepts (line 462) | def _p_get_has_nsfw_concepts(pipe, features, params):
function unshard (line 466) | def unshard(x: jnp.ndarray):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py
class FlaxStableDiffusionImg2ImgPipeline (line 105) | class FlaxStableDiffusionImg2ImgPipeline(FlaxDiffusionPipeline):
method __init__ (line 134) | def __init__(
method prepare_inputs (line 171) | def prepare_inputs(self, prompt: Union[str, List[str]], image: Union[I...
method _get_has_nsfw_concepts (line 192) | def _get_has_nsfw_concepts(self, features, params):
method _run_safety_checker (line 196) | def _run_safety_checker(self, images, safety_model_params, jit=False):
method get_timestep_start (line 226) | def get_timestep_start(self, num_inference_steps, strength):
method _generate (line 234) | def _generate(
method __call__ (line 339) | def __call__(
function _p_generate (line 480) | def _p_generate(
function _p_get_has_nsfw_concepts (line 510) | def _p_get_has_nsfw_concepts(pipe, features, params):
function unshard (line 514) | def unshard(x: jnp.ndarray):
function preprocess (line 521) | def preprocess(image, dtype):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py
class FlaxStableDiffusionInpaintPipeline (line 102) | class FlaxStableDiffusionInpaintPipeline(FlaxDiffusionPipeline):
method __init__ (line 131) | def __init__(
method prepare_inputs (line 189) | def prepare_inputs(
method _get_has_nsfw_concepts (line 228) | def _get_has_nsfw_concepts(self, features, params):
method _run_safety_checker (line 232) | def _run_safety_checker(self, images, safety_model_params, jit=False):
method _generate (line 262) | def _generate(
method __call__ (line 390) | def __call__(
function _p_generate (line 523) | def _p_generate(
function _p_get_has_nsfw_concepts (line 553) | def _p_get_has_nsfw_concepts(pipe, features, params):
function unshard (line 557) | def unshard(x: jnp.ndarray):
function preprocess_image (line 564) | def preprocess_image(image, dtype):
function preprocess_mask (line 573) | def preprocess_mask(mask, dtype):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py
class OnnxStableDiffusionPipeline (line 33) | class OnnxStableDiffusionPipeline(DiffusionPipeline):
method __init__ (line 45) | def __init__(
method _encode_prompt (line 114) | def _encode_prompt(
method check_inputs (line 217) | def check_inputs(
method __call__ (line 264) | def __call__(
class StableDiffusionOnnxPipeline (line 462) | class StableDiffusionOnnxPipeline(OnnxStableDiffusionPipeline):
method __init__ (line 463) | def __init__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py
function preprocess (line 36) | def preprocess(image):
class OnnxStableDiffusionImg2ImgPipeline (line 64) | class OnnxStableDiffusionImg2ImgPipeline(DiffusionPipeline):
method __init__ (line 102) | def __init__(
method _encode_prompt (line 172) | def _encode_prompt(
method check_inputs (line 275) | def check_inputs(
method __call__ (line 317) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py
function prepare_mask_and_masked_image (line 38) | def prepare_mask_and_masked_image(image, mask, latents_shape):
class OnnxStableDiffusionInpaintPipeline (line 56) | class OnnxStableDiffusionInpaintPipeline(DiffusionPipeline):
method __init__ (line 94) | def __init__(
method _encode_prompt (line 165) | def _encode_prompt(
method check_inputs (line 269) | def check_inputs(
method __call__ (line 317) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint_legacy.py
function preprocess (line 20) | def preprocess(image):
function preprocess_mask (line 29) | def preprocess_mask(mask, scale_factor=8):
class OnnxStableDiffusionInpaintPipelineLegacy (line 41) | class OnnxStableDiffusionInpaintPipelineLegacy(DiffusionPipeline):
method __init__ (line 80) | def __init__(
method _encode_prompt (line 150) | def _encode_prompt(
method check_inputs (line 253) | def check_inputs(
method __call__ (line 295) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py
function preprocess (line 26) | def preprocess(image):
class OnnxStableDiffusionUpscalePipeline (line 48) | class OnnxStableDiffusionUpscalePipeline(StableDiffusionUpscalePipeline):
method __init__ (line 49) | def __init__(
method __call__ (line 72) | def __call__(
method decode_latents (line 281) | def decode_latents(self, latents):
method _encode_prompt (line 288) | def _encode_prompt(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
function rescale_noise_cfg (line 58) | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
class StableDiffusionPipeline (line 72) | class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderM...
method __init__ (line 109) | def __init__(
method enable_vae_slicing (line 199) | def enable_vae_slicing(self):
method disable_vae_slicing (line 208) | def disable_vae_slicing(self):
method enable_vae_tiling (line 215) | def enable_vae_tiling(self):
method disable_vae_tiling (line 224) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 231) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 256) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 285) | def _execution_device(self):
method _encode_prompt (line 302) | def _encode_prompt(
method run_safety_checker (line 456) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 470) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 483) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 500) | def check_inputs(
method prepare_latents (line 547) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 566) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py
class AttentionStore (line 74) | class AttentionStore:
method get_empty_store (line 76) | def get_empty_store():
method __call__ (line 79) | def __call__(self, attn, is_cross: bool, place_in_unet: str):
method between_steps (line 89) | def between_steps(self):
method get_average_attention (line 93) | def get_average_attention(self):
method aggregate_attention (line 97) | def aggregate_attention(self, from_where: List[str]) -> torch.Tensor:
method reset (line 109) | def reset(self):
method __init__ (line 114) | def __init__(self, attn_res):
class AttendExciteAttnProcessor (line 127) | class AttendExciteAttnProcessor:
method __init__ (line 128) | def __init__(self, attnstore, place_in_unet):
method __call__ (line 133) | def __call__(self, attn: Attention, hidden_states, encoder_hidden_stat...
class StableDiffusionAttendAndExcitePipeline (line 165) | class StableDiffusionAttendAndExcitePipeline(DiffusionPipeline, TextualI...
method __init__ (line 194) | def __init__(
method enable_vae_slicing (line 237) | def enable_vae_slicing(self):
method disable_vae_slicing (line 247) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 255) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 282) | def _execution_device(self):
method _encode_prompt (line 300) | def _encode_prompt(
method run_safety_checker (line 455) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 470) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 484) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 501) | def check_inputs(
method prepare_latents (line 575) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method _compute_max_attention_per_index (line 593) | def _compute_max_attention_per_index(
method _aggregate_and_get_max_attention_per_token (line 615) | def _aggregate_and_get_max_attention_per_token(
method _compute_loss (line 630) | def _compute_loss(max_attention_per_index: List[torch.Tensor]) -> torc...
method _update_latent (line 637) | def _update_latent(latents: torch.Tensor, loss: torch.Tensor, step_siz...
method _perform_iterative_refinement_step (line 643) | def _perform_iterative_refinement_step(
method register_attention_control (line 697) | def register_attention_control(self):
method get_indices (line 716) | def get_indices(self, prompt: str) -> Dict[str, int]:
method __call__ (line 724) | def __call__(
class GaussianSmoothing (line 1015) | class GaussianSmoothing(torch.nn.Module):
method __init__ (line 1028) | def __init__(
method forward (line 1069) | def forward(self, input):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py
function preprocess (line 39) | def preprocess(image):
class StableDiffusionDepth2ImgPipeline (line 65) | class StableDiffusionDepth2ImgPipeline(DiffusionPipeline, TextualInversi...
method __init__ (line 95) | def __init__(
method enable_sequential_cpu_offload (line 140) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 159) | def _execution_device(self):
method _encode_prompt (line 177) | def _encode_prompt(
method run_safety_checker (line 332) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 347) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 361) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 379) | def check_inputs(
method get_timesteps (line 420) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 430) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method prepare_depth_map (line 487) | def prepare_depth_map(self, image, depth_map, batch_size, do_classifie...
method __call__ (line 532) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_diffedit.py
class DiffEditInversionPipelineOutput (line 50) | class DiffEditInversionPipelineOutput(BaseOutput):
function auto_corr_loss (line 140) | def auto_corr_loss(hidden_states, generator=None):
function kl_divergence (line 156) | def kl_divergence(hidden_states):
function preprocess (line 161) | def preprocess(image):
function preprocess_mask (line 187) | def preprocess_mask(mask, batch_size: int = 1):
class StableDiffusionDiffEditPipeline (line 240) | class StableDiffusionDiffEditPipeline(DiffusionPipeline, TextualInversio...
method __init__ (line 277) | def __init__(
method enable_vae_slicing (line 371) | def enable_vae_slicing(self):
method disable_vae_slicing (line 381) | def disable_vae_slicing(self):
method enable_vae_tiling (line 389) | def enable_vae_tiling(self):
method disable_vae_tiling (line 399) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 407) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 433) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 463) | def _execution_device(self):
method _encode_prompt (line 481) | def _encode_prompt(
method run_safety_checker (line 636) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 651) | def prepare_extra_step_kwargs(self, generator, eta):
method decode_latents (line 669) | def decode_latents(self, latents):
method check_inputs (line 682) | def check_inputs(
method check_source_inputs (line 730) | def check_source_inputs(
method get_timesteps (line 766) | def get_timesteps(self, num_inference_steps, strength, device):
method get_inverse_timesteps (line 775) | def get_inverse_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 789) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method prepare_image_latents (line 807) | def prepare_image_latents(self, image, batch_size, dtype, device, gene...
method get_epsilon (line 856) | def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor...
method generate_mask (line 874) | def generate_mask(
method invert (line 1097) | def invert(
method __call__ (line 1335) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
class StableDiffusionImageVariationPipeline (line 37) | class StableDiffusionImageVariationPipeline(DiffusionPipeline):
method __init__ (line 65) | def __init__(
method enable_sequential_cpu_offload (line 126) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 145) | def _execution_device(self):
method _encode_image (line 162) | def _encode_image(self, image, device, num_images_per_prompt, do_class...
method run_safety_checker (line 188) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 203) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 217) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 234) | def check_inputs(self, image, height, width, callback_steps):
method prepare_latents (line 257) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 275) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
function preprocess (line 75) | def preprocess(image):
class StableDiffusionImg2ImgPipeline (line 101) | class StableDiffusionImg2ImgPipeline(
method __init__ (line 140) | def __init__(
method enable_sequential_cpu_offload (line 231) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 257) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 287) | def _execution_device(self):
method _encode_prompt (line 305) | def _encode_prompt(
method run_safety_checker (line 460) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 475) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 489) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 506) | def check_inputs(
method get_timesteps (line 546) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 555) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method __call__ (line 614) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
function prepare_mask_and_masked_image (line 39) | def prepare_mask_and_masked_image(image, mask, height, width, return_ima...
class StableDiffusionInpaintPipeline (line 156) | class StableDiffusionInpaintPipeline(
method __init__ (line 204) | def __init__(
method enable_sequential_cpu_offload (line 300) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 326) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 356) | def _execution_device(self):
method _encode_prompt (line 374) | def _encode_prompt(
method run_safety_checker (line 529) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 544) | def prepare_extra_step_kwargs(self, generator, eta):
method decode_latents (line 562) | def decode_latents(self, latents):
method check_inputs (line 575) | def check_inputs(
method prepare_latents (line 626) | def prepare_latents(
method _encode_vae_image (line 679) | def _encode_vae_image(self, image: torch.Tensor, generator: torch.Gene...
method prepare_mask_latents (line 693) | def prepare_mask_latents(
method get_timesteps (line 735) | def get_timesteps(self, num_inference_steps, strength, device):
method __call__ (line 745) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py
function preprocess_image (line 46) | def preprocess_image(image, batch_size):
function preprocess_mask (line 56) | def preprocess_mask(mask, batch_size, scale_factor=8):
class StableDiffusionInpaintPipelineLegacy (line 87) | class StableDiffusionInpaintPipelineLegacy(
method __init__ (line 126) | def __init__(
method enable_sequential_cpu_offload (line 224) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 250) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 280) | def _execution_device(self):
method _encode_prompt (line 298) | def _encode_prompt(
method run_safety_checker (line 453) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 468) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 482) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 500) | def check_inputs(
method get_timesteps (line 541) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 550) | def prepare_latents(self, image, timestep, num_images_per_prompt, dtyp...
method __call__ (line 567) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
function preprocess (line 45) | def preprocess(image):
class StableDiffusionInstructPix2PixPipeline (line 71) | class StableDiffusionInstructPix2PixPipeline(DiffusionPipeline, TextualI...
method __init__ (line 107) | def __init__(
method __call__ (line 150) | def __call__(
method enable_sequential_cpu_offload (line 432) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 458) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 488) | def _execution_device(self):
method _encode_prompt (line 505) | def _encode_prompt(
method run_safety_checker (line 653) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 668) | def prepare_extra_step_kwargs(self, generator, eta):
method decode_latents (line 686) | def decode_latents(self, latents):
method check_inputs (line 699) | def check_inputs(
method prepare_latents (line 737) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method prepare_image_latents (line 754) | def prepare_image_latents(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py
class ModelWrapper (line 34) | class ModelWrapper:
method __init__ (line 35) | def __init__(self, model, alphas_cumprod):
method apply_model (line 39) | def apply_model(self, *args, **kwargs):
class StableDiffusionKDiffusionPipeline (line 48) | class StableDiffusionKDiffusionPipeline(DiffusionPipeline, TextualInvers...
method __init__ (line 83) | def __init__(
method set_scheduler (line 124) | def set_scheduler(self, scheduler_type: str):
method enable_sequential_cpu_offload (line 130) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 156) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 186) | def _execution_device(self):
method _encode_prompt (line 204) | def _encode_prompt(
method run_safety_checker (line 359) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 374) | def decode_latents(self, latents):
method check_inputs (line 388) | def check_inputs(
method prepare_latents (line 435) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 448) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py
function preprocess (line 35) | def preprocess(image):
class StableDiffusionLatentUpscalePipeline (line 61) | class StableDiffusionLatentUpscalePipeline(DiffusionPipeline):
method __init__ (line 84) | def __init__(
method enable_sequential_cpu_offload (line 104) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 123) | def _execution_device(self):
method _encode_prompt (line 140) | def _encode_prompt(self, prompt, device, do_classifier_free_guidance, ...
method decode_latents (line 231) | def decode_latents(self, latents):
method check_inputs (line 244) | def check_inputs(self, prompt, image, callback_steps):
method prepare_latents (line 282) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 296) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py
class LDM3DPipelineOutput (line 59) | class LDM3DPipelineOutput(BaseOutput):
class StableDiffusionLDM3DPipeline (line 77) | class StableDiffusionLDM3DPipeline(
method __init__ (line 118) | def __init__(
method enable_vae_slicing (line 161) | def enable_vae_slicing(self):
method disable_vae_slicing (line 171) | def disable_vae_slicing(self):
method enable_vae_tiling (line 179) | def enable_vae_tiling(self):
method disable_vae_tiling (line 189) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 196) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 222) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 252) | def _execution_device(self):
method _encode_prompt (line 270) | def _encode_prompt(
method run_safety_checker (line 424) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 440) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 458) | def check_inputs(
method prepare_latents (line 505) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 524) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_model_editing.py
class StableDiffusionModelEditingPipeline (line 58) | class StableDiffusionModelEditingPipeline(DiffusionPipeline, TextualInve...
method __init__ (line 90) | def __init__(
method enable_vae_slicing (line 168) | def enable_vae_slicing(self):
method disable_vae_slicing (line 178) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 186) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 213) | def _execution_device(self):
method _encode_prompt (line 231) | def _encode_prompt(
method run_safety_checker (line 386) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 401) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 415) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 433) | def check_inputs(
method prepare_latents (line 481) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method edit_model (line 499) | def edit_model(
method __call__ (line 617) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_panorama.py
class StableDiffusionPanoramaPipeline (line 54) | class StableDiffusionPanoramaPipeline(DiffusionPipeline, TextualInversio...
method __init__ (line 87) | def __init__(
method enable_vae_slicing (line 130) | def enable_vae_slicing(self):
method disable_vae_slicing (line 140) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 148) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 175) | def _execution_device(self):
method _encode_prompt (line 193) | def _encode_prompt(
method run_safety_checker (line 348) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 363) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 377) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 395) | def check_inputs(
method prepare_latents (line 443) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method get_views (line 460) | def get_views(self, panorama_height, panorama_width, window_size=64, s...
method __call__ (line 479) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_paradigms.py
class StableDiffusionParadigmsPipeline (line 62) | class StableDiffusionParadigmsPipeline(
method __init__ (line 104) | def __init__(
method enable_vae_slicing (line 150) | def enable_vae_slicing(self):
method disable_vae_slicing (line 160) | def disable_vae_slicing(self):
method enable_vae_tiling (line 168) | def enable_vae_tiling(self):
method disable_vae_tiling (line 178) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 186) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 212) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 242) | def _execution_device(self):
method _encode_prompt (line 260) | def _encode_prompt(
method run_safety_checker (line 415) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 430) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 448) | def check_inputs(
method prepare_latents (line 496) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method _cumsum (line 513) | def _cumsum(self, input, dim, debug=False):
method __call__ (line 523) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py
class Pix2PixInversionPipelineOutput (line 57) | class Pix2PixInversionPipelineOutput(BaseOutput, TextualInversionLoaderM...
function preprocess (line 179) | def preprocess(image):
function prepare_unet (line 205) | def prepare_unet(unet: UNet2DConditionModel):
class Pix2PixZeroL2Loss (line 222) | class Pix2PixZeroL2Loss:
method __init__ (line 223) | def __init__(self):
method compute_loss (line 226) | def compute_loss(self, predictions, targets):
class Pix2PixZeroAttnProcessor (line 230) | class Pix2PixZeroAttnProcessor:
method __init__ (line 234) | def __init__(self, is_pix2pix_zero=False):
method __call__ (line 239) | def __call__(
class StableDiffusionPix2PixZeroPipeline (line 285) | class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline):
method __init__ (line 323) | def __init__(
method enable_sequential_cpu_offload (line 372) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 397) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 423) | def _execution_device(self):
method _encode_prompt (line 441) | def _encode_prompt(
method run_safety_checker (line 596) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 611) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 625) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 642) | def check_inputs(
method prepare_latents (line 673) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method generate_caption (line 691) | def generate_caption(self, images):
method construct_direction (line 710) | def construct_direction(self, embs_source: torch.Tensor, embs_target: ...
method get_embeds (line 715) | def get_embeds(self, prompt: List[str], batch_size: int = 16) -> torch...
method prepare_image_latents (line 734) | def prepare_image_latents(self, image, batch_size, dtype, device, gene...
method get_epsilon (line 783) | def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor...
method auto_corr_loss (line 800) | def auto_corr_loss(self, hidden_states, generator=None):
method kl_divergence (line 815) | def kl_divergence(self, hidden_states):
method __call__ (line 822) | def __call__(
method invert (line 1101) | def invert(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py
class CrossAttnStoreProcessor (line 53) | class CrossAttnStoreProcessor:
method __init__ (line 54) | def __init__(self):
method __call__ (line 57) | def __call__(
class StableDiffusionSAGPipeline (line 93) | class StableDiffusionSAGPipeline(DiffusionPipeline, TextualInversionLoad...
method __init__ (line 122) | def __init__(
method enable_vae_slicing (line 149) | def enable_vae_slicing(self):
method disable_vae_slicing (line 159) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 167) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 194) | def _execution_device(self):
method _encode_prompt (line 212) | def _encode_prompt(
method run_safety_checker (line 367) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 382) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 396) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 414) | def check_inputs(
method prepare_latents (line 462) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 481) | def __call__(
method sag_masking (line 724) | def sag_masking(self, original_latents, attn_map, map_size, t, eps):
method pred_x0 (line 754) | def pred_x0(self, sample, model_output, timestep):
method pred_epsilon (line 774) | def pred_epsilon(self, sample, model_output, timestep):
function gaussian_blur_2d (line 794) | def gaussian_blur_2d(img, kernel_size, sigma):
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
function preprocess (line 42) | def preprocess(image):
class StableDiffusionUpscalePipeline (line 68) | class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversion...
method __init__ (line 95) | def __init__(
method enable_sequential_cpu_offload (line 143) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 160) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 188) | def _execution_device(self):
method run_safety_checker (line 206) | def run_safety_checker(self, image, device, dtype):
method _encode_prompt (line 223) | def _encode_prompt(
method prepare_extra_step_kwargs (line 378) | def prepare_extra_step_kwargs(self, generator, eta):
method decode_latents (line 396) | def decode_latents(self, latents):
method check_inputs (line 409) | def check_inputs(
method prepare_latents (line 491) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 505) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
class StableUnCLIPPipeline (line 53) | class StableUnCLIPPipeline(DiffusionPipeline, TextualInversionLoaderMixi...
method __init__ (line 106) | def __init__(
method enable_vae_slicing (line 144) | def enable_vae_slicing(self):
method disable_vae_slicing (line 154) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 161) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 185) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 212) | def _execution_device(self):
method _encode_prior_prompt (line 230) | def _encode_prior_prompt(
method _encode_prompt (line 332) | def _encode_prompt(
method decode_latents (line 487) | def decode_latents(self, latents):
method prepare_prior_extra_step_kwargs (line 501) | def prepare_prior_extra_step_kwargs(self, generator, eta):
method prepare_extra_step_kwargs (line 519) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 536) | def check_inputs(
method prepare_latents (line 597) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method noise_image_embeddings (line 608) | def noise_image_embeddings(
method __call__ (line 656) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
class StableUnCLIPImg2ImgPipeline (line 66) | class StableUnCLIPImg2ImgPipeline(DiffusionPipeline, TextualInversionLoa...
method __init__ (line 112) | def __init__(
method enable_vae_slicing (line 146) | def enable_vae_slicing(self):
method disable_vae_slicing (line 156) | def disable_vae_slicing(self):
method enable_sequential_cpu_offload (line 163) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 187) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 214) | def _execution_device(self):
method _encode_prompt (line 232) | def _encode_prompt(
method _encode_image (line 386) | def _encode_image(
method decode_latents (line 442) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 456) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 473) | def check_inputs(
method prepare_latents (line 557) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method noise_image_embeddings (line 575) | def noise_image_embeddings(
method __call__ (line 623) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/safety_checker.py
function cosine_distance (line 26) | def cosine_distance(image_embeds, text_embeds):
class StableDiffusionSafetyChecker (line 32) | class StableDiffusionSafetyChecker(PreTrainedModel):
method __init__ (line 37) | def __init__(self, config: CLIPConfig):
method forward (line 50) | def forward(self, clip_input, images):
method forward_onnx (line 102) | def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.Fl...
FILE: 4DoF/diffusers/pipelines/stable_diffusion/safety_checker_flax.py
function jax_cosine_distance (line 25) | def jax_cosine_distance(emb_1, emb_2, eps=1e-12):
class FlaxStableDiffusionSafetyCheckerModule (line 31) | class FlaxStableDiffusionSafetyCheckerModule(nn.Module):
method setup (line 35) | def setup(self):
method __call__ (line 47) | def __call__(self, clip_input):
class FlaxStableDiffusionSafetyChecker (line 71) | class FlaxStableDiffusionSafetyChecker(FlaxPreTrainedModel):
method __init__ (line 76) | def __init__(
method init_weights (line 90) | def init_weights(self, rng: jax.random.KeyArray, input_shape: Tuple, p...
method __call__ (line 101) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py
class StableUnCLIPImageNormalizer (line 24) | class StableUnCLIPImageNormalizer(ModelMixin, ConfigMixin):
method __init__ (line 33) | def __init__(
method to (line 42) | def to(
method scale (line 51) | def scale(self, embeds):
method unscale (line 55) | def unscale(self, embeds):
FILE: 4DoF/diffusers/pipelines/stable_diffusion_safe/__init__.py
class SafetyConfig (line 13) | class SafetyConfig(object):
class StableDiffusionSafePipelineOutput (line 45) | class StableDiffusionSafePipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py
class StableDiffusionPipelineSafe (line 22) | class StableDiffusionPipelineSafe(DiffusionPipeline):
method __init__ (line 54) | def __init__(
method safety_concept (line 150) | def safety_concept(self):
method safety_concept (line 160) | def safety_concept(self, concept):
method enable_sequential_cpu_offload (line 170) | def enable_sequential_cpu_offload(self):
method _execution_device (line 189) | def _execution_device(self):
method _encode_prompt (line 206) | def _encode_prompt(
method run_safety_checker (line 341) | def run_safety_checker(self, image, device, dtype, enable_safety_guida...
method decode_latents (line 365) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 379) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 397) | def check_inputs(
method prepare_latents (line 445) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method perform_safety_guidance (line 462) | def perform_safety_guidance(
method __call__ (line 505) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion_safe/safety_checker.py
function cosine_distance (line 25) | def cosine_distance(image_embeds, text_embeds):
class SafeStableDiffusionSafetyChecker (line 31) | class SafeStableDiffusionSafetyChecker(PreTrainedModel):
method __init__ (line 36) | def __init__(self, config: CLIPConfig):
method forward (line 49) | def forward(self, clip_input, images):
method forward_onnx (line 88) | def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.Fl...
FILE: 4DoF/diffusers/pipelines/stable_diffusion_xl/__init__.py
class StableDiffusionXLPipelineOutput (line 11) | class StableDiffusionXLPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
function rescale_noise_cfg (line 62) | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
class StableDiffusionXLPipeline (line 76) | class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin):
method __init__ (line 107) | def __init__(
method enable_vae_slicing (line 137) | def enable_vae_slicing(self):
method disable_vae_slicing (line 147) | def disable_vae_slicing(self):
method enable_vae_tiling (line 155) | def enable_vae_tiling(self):
method disable_vae_tiling (line 165) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 172) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 194) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 226) | def _execution_device(self):
method encode_prompt (line 243) | def encode_prompt(
method prepare_extra_step_kwargs (line 430) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 447) | def check_inputs(
method prepare_latents (line 507) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method _get_add_time_ids (line 524) | def _get_add_time_ids(self, original_size, crops_coords_top_left, targ...
method __call__ (line 542) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
function rescale_noise_cfg (line 67) | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
class StableDiffusionXLImg2ImgPipeline (line 81) | class StableDiffusionXLImg2ImgPipeline(DiffusionPipeline, FromSingleFile...
method __init__ (line 113) | def __init__(
method enable_vae_slicing (line 144) | def enable_vae_slicing(self):
method disable_vae_slicing (line 154) | def disable_vae_slicing(self):
method enable_vae_tiling (line 162) | def enable_vae_tiling(self):
method disable_vae_tiling (line 172) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 180) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 203) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 235) | def _execution_device(self):
method encode_prompt (line 253) | def encode_prompt(
method prepare_extra_step_kwargs (line 440) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 457) | def check_inputs(
method get_timesteps (line 497) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 506) | def prepare_latents(self, image, timestep, batch_size, num_images_per_...
method _get_add_time_ids (line 568) | def _get_add_time_ids(
method __call__ (line 609) | def __call__(
FILE: 4DoF/diffusers/pipelines/stable_diffusion_xl/watermark.py
class StableDiffusionXLWatermarker (line 12) | class StableDiffusionXLWatermarker:
method __init__ (line 13) | def __init__(self):
method apply_watermark (line 19) | def apply_watermark(self, images: torch.FloatTensor):
FILE: 4DoF/diffusers/pipelines/stochastic_karras_ve/pipeline_stochastic_karras_ve.py
class KarrasVePipeline (line 25) | class KarrasVePipeline(DiffusionPipeline):
method __init__ (line 44) | def __init__(self, unet: UNet2DModel, scheduler: KarrasVeScheduler):
method __call__ (line 49) | def __call__(
FILE: 4DoF/diffusers/pipelines/text_to_video_synthesis/__init__.py
class TextToVideoSDPipelineOutput (line 11) | class TextToVideoSDPipelineOutput(BaseOutput):
FILE: 4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py
function tensor2vid (line 58) | def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5,...
class TextToVideoSDPipeline (line 76) | class TextToVideoSDPipeline(DiffusionPipeline, TextualInversionLoaderMix...
method __init__ (line 97) | def __init__(
method enable_vae_slicing (line 117) | def enable_vae_slicing(self):
method disable_vae_slicing (line 127) | def disable_vae_slicing(self):
method enable_vae_tiling (line 135) | def enable_vae_tiling(self):
method disable_vae_tiling (line 145) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 152) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 173) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 200) | def _execution_device(self):
method _encode_prompt (line 218) | def _encode_prompt(
method decode_latents (line 372) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 396) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 414) | def check_inputs(
method prepare_latents (line 461) | def prepare_latents(
method __call__ (line 488) | def __call__(
FILE: 4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
function tensor2vid (line 79) | def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5,...
function preprocess_video (line 97) | def preprocess_video(video):
class VideoToVideoSDPipeline (line 138) | class VideoToVideoSDPipeline(DiffusionPipeline, TextualInversionLoaderMi...
method __init__ (line 159) | def __init__(
method enable_vae_slicing (line 179) | def enable_vae_slicing(self):
method disable_vae_slicing (line 189) | def disable_vae_slicing(self):
method enable_vae_tiling (line 197) | def enable_vae_tiling(self):
method disable_vae_tiling (line 207) | def disable_vae_tiling(self):
method enable_sequential_cpu_offload (line 214) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 235) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 262) | def _execution_device(self):
method _encode_prompt (line 280) | def _encode_prompt(
method decode_latents (line 435) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 459) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 477) | def check_inputs(
method get_timesteps (line 518) | def get_timesteps(self, num_inference_steps, strength, device):
method prepare_latents (line 527) | def prepare_latents(self, video, timestep, batch_size, dtype, device, ...
method __call__ (line 573) | def __call__(
FILE: 4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py
function rearrange_0 (line 18) | def rearrange_0(tensor, f):
function rearrange_1 (line 24) | def rearrange_1(tensor):
function rearrange_3 (line 29) | def rearrange_3(tensor, f):
function rearrange_4 (line 34) | def rearrange_4(tensor):
class CrossFrameAttnProcessor (line 39) | class CrossFrameAttnProcessor:
method __init__ (line 49) | def __init__(self, batch_size=2):
method __call__ (line 52) | def __call__(self, attn, hidden_states, encoder_hidden_states=None, at...
class CrossFrameAttnProcessor2_0 (line 98) | class CrossFrameAttnProcessor2_0:
method __init__ (line 108) | def __init__(self, batch_size=2):
method __call__ (line 113) | def __call__(self, attn, hidden_states, encoder_hidden_states=None, at...
class TextToVideoPipelineOutput (line 174) | class TextToVideoPipelineOutput(BaseOutput):
function coords_grid (line 179) | def coords_grid(batch, ht, wd, device):
function warp_single_latent (line 186) | def warp_single_latent(latent, reference_flow):
function create_motion_field (line 213) | def create_motion_field(motion_field_strength_x, motion_field_strength_y...
function create_motion_field_and_warp_latents (line 236) | def create_motion_field_and_warp_latents(motion_field_strength_x, motion...
class TextToVideoZeroPipeline (line 263) | class TextToVideoZeroPipeline(StableDiffusionPipeline):
method __init__ (line 291) | def __init__(
method forward_loop (line 312) | def forward_loop(self, x_t0, t0, t1, generator):
method backward_loop (line 330) | def backward_loop(
method __call__ (line 400) | def __call__(
FILE: 4DoF/diffusers/pipelines/unclip/pipeline_unclip.py
class UnCLIPPipeline (line 34) | class UnCLIPPipeline(DiffusionPipeline):
method __init__ (line 78) | def __init__(
method prepare_latents (line 106) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method _encode_prompt (line 117) | def _encode_prompt(
method enable_sequential_cpu_offload (line 208) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 234) | def _execution_device(self):
method __call__ (line 252) | def __call__(
FILE: 4DoF/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py
class UnCLIPImageVariationPipeline (line 38) | class UnCLIPImageVariationPipeline(DiffusionPipeline):
method __init__ (line 84) | def __init__(
method prepare_latents (line 113) | def prepare_latents(self, shape, dtype, device, generator, latents, sc...
method _encode_prompt (line 124) | def _encode_prompt(self, prompt, device, num_images_per_prompt, do_cla...
method _encode_image (line 187) | def _encode_image(self, image, device, num_images_per_prompt, image_em...
method enable_sequential_cpu_offload (line 201) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 227) | def _execution_device(self):
method __call__ (line 245) | def __call__(
FILE: 4DoF/diffusers/pipelines/unclip/text_proj.py
class UnCLIPTextProjModel (line 22) | class UnCLIPTextProjModel(ModelMixin, ConfigMixin):
method __init__ (line 31) | def __init__(
method forward (line 55) | def forward(self, *, image_embeddings, prompt_embeds, text_encoder_hid...
FILE: 4DoF/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
class UniDiffuserTextDecoder (line 14) | class UniDiffuserTextDecoder(ModelMixin, ConfigMixin, ModuleUtilsMixin):
method __init__ (line 67) | def __init__(
method forward (line 131) | def forward(
method get_dummy_token (line 163) | def get_dummy_token(self, batch_size: int, device: torch.device) -> to...
method encode (line 166) | def encode(self, prefix):
method generate_captions (line 170) | def generate_captions(self, features, eos_token_id, device):
method generate_beam (line 202) | def generate_beam(
FILE: 4DoF/diffusers/pipelines/unidiffuser/modeling_uvit.py
function _no_grad_trunc_normal_ (line 19) | def _no_grad_trunc_normal_(tensor, mean, std, a, b):
function trunc_normal_ (line 56) | def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
class PatchEmbed (line 75) | class PatchEmbed(nn.Module):
method __init__ (line 78) | def __init__(
method forward (line 109) | def forward(self, latent):
class SkipBlock (line 121) | class SkipBlock(nn.Module):
method __init__ (line 122) | def __init__(self, dim: int):
method forward (line 130) | def forward(self, x, skip):
class UTransformerBlock (line 140) | class UTransformerBlock(nn.Module):
method __init__ (line 174) | def __init__(
method forward (line 251) | def forward(
class UniDiffuserBlock (line 333) | class UniDiffuserBlock(nn.Module):
method __init__ (line 369) | def __init__(
method forward (line 446) | def forward(
class UTransformer2DModel (line 530) | class UTransformer2DModel(ModelMixin, ConfigMixin):
method __init__ (line 589) | def __init__(
method forward (line 727) | def forward(
class UniDiffuserModel (line 832) | class UniDiffuserModel(ModelMixin, ConfigMixin):
method __init__ (line 897) | def __init__(
method no_weight_decay (line 1034) | def no_weight_decay(self):
method forward (line 1037) | def forward(
FILE: 4DoF/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
function preprocess (line 37) | def preprocess(image):
class ImageTextPipelineOutput (line 65) | class ImageTextPipelineOutput(BaseOutput):
class UniDiffuserPipeline (line 82) | class UniDiffuserPipeline(DiffusionPipeline):
method __init__ (line 126) | def __init__(
method enable_sequential_cpu_offload (line 177) | def enable_sequential_cpu_offload(self, gpu_id=0):
method enable_model_cpu_offload (line 204) | def enable_model_cpu_offload(self, gpu_id=0):
method _execution_device (line 234) | def _execution_device(self):
method prepare_extra_step_kwargs (line 252) | def prepare_extra_step_kwargs(self, generator, eta):
method _infer_mode (line 269) | def _infer_mode(self, prompt, prompt_embeds, image, latents, prompt_la...
method set_text_mode (line 328) | def set_text_mode(self):
method set_image_mode (line 332) | def set_image_mode(self):
method set_text_to_image_mode (line 336) | def set_text_to_image_mode(self):
method set_image_to_text_mode (line 340) | def set_image_to_text_mode(self):
method set_joint_mode (line 344) | def set_joint_mode(self):
method reset_mode (line 348) | def reset_mode(self):
method _infer_batch_size (line 352) | def _infer_batch_size(
method _encode_prompt (line 430) | def _encode_prompt(
method encode_image_vae_latents (line 570) | def encode_image_vae_latents(
method encode_image_clip_latents (line 630) | def encode_image_clip_latents(
method decode_image_latents (line 689) | def decode_image_latents(self, latents):
method prepare_text_latents (line 697) | def prepare_text_latents(
method prepare_image_vae_latents (line 721) | def prepare_image_vae_latents(
method prepare_image_clip_latents (line 756) | def prepare_image_clip_latents(
method _split (line 778) | def _split(self, x, height, width):
method _combine (line 794) | def _combine(self, img_vae, img_clip):
method _split_joint (line 803) | def _split_joint(self, x, height, width):
method _combine_joint (line 822) | def _combine_joint(self, img_vae, img_clip, text):
method _get_noise_pred (line 833) | def _get_noise_pred(
method check_latents_shape (line 955) | def check_latents_shape(self, latents_name, latents, expected_shape):
method check_inputs (line 971) | def check_inputs(
method __call__ (line 1085) | def __call__(
FILE: 4DoF/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py
function get_down_block (line 38) | def get_down_block(
function get_up_block (line 97) | def get_up_block(
class UNetFlatConditionModel (line 157) | class UNetFlatConditionModel(ModelMixin, ConfigMixin):
method __init__ (line 248) | def __init__(
method attn_processors (line 674) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 697) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 731) | def set_default_attn_processor(self):
method set_attention_slice (line 737) | def set_attention_slice(self, slice_size):
method _set_gradient_checkpointing (line 802) | def _set_gradient_checkpointing(self, module, value=False):
method forward (line 806) | def forward(
class LinearMultiDim (line 1096) | class LinearMultiDim(nn.Linear):
method __init__ (line 1097) | def __init__(self, in_features, out_features=None, second_dim=4, *args...
method forward (line 1106) | def forward(self, input_tensor, *args, **kwargs):
class ResnetBlockFlat (line 1115) | class ResnetBlockFlat(nn.Module):
method __init__ (line 1116) | def __init__(
method forward (line 1176) | def forward(self, input_tensor, temb):
class DownBlockFlat (line 1210) | class DownBlockFlat(nn.Module):
method __init__ (line 1211) | def __init__(
method forward (line 1262) | def forward(self, hidden_states, temb=None):
class CrossAttnDownBlockFlat (line 1297) | class CrossAttnDownBlockFlat(nn.Module):
method __init__ (line 1298) | def __init__(
method forward (line 1385) | def forward(
class UpBlockFlat (line 1449) | class UpBlockFlat(nn.Module):
method __init__ (line 1450) | def __init__(
method forward (line 1497) | def forward(self, hidden_states, res_hidden_states_tuple, temb=None, u...
class CrossAttnUpBlockFlat (line 1531) | class CrossAttnUpBlockFlat(nn.Module):
method __init__ (line 1532) | def __init__(
method forward (line 1615) | def forward(
class UNetMidBlockFlatCrossAttn (line 1680) | class UNetMidBlockFlatCrossAttn(nn.Module):
method __init__ (line 1681) | def __init__(
method forward (line 1766) | def forward(
class UNetMidBlockFlatSimpleCrossAttn (line 1791) | class UNetMidBlockFlatSimpleCrossAttn(nn.Module):
method __init__ (line 1792) | def __init__(
method forward (line 1876) | def forward(
FILE: 4DoF/diffusers/pipelines/versatile_diffusion/pipeline_versatile_diffusion.py
class VersatileDiffusionPipeline (line 20) | class VersatileDiffusionPipeline(DiffusionPipeline):
method __init__ (line 57) | def __init__(
method image_variation (line 83) | def image_variation(
method text_to_image (line 199) | def text_to_image(
method dual_guided (line 311) | def dual_guided(
FILE: 4DoF/diffusers/pipelines/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py
class VersatileDiffusionDualGuidedPipeline (line 41) | class VersatileDiffusionDualGuidedPipeline(DiffusionPipeline):
method __init__ (line 70) | def __init__(
method remove_unused_weights (line 101) | def remove_unused_weights(self):
method _convert_to_dual_attention (line 104) | def _convert_to_dual_attention(self):
method _revert_dual_attention (line 138) | def _revert_dual_attention(self):
method enable_sequential_cpu_offload (line 151) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 170) | def _execution_device(self):
method _encode_text_prompt (line 187) | def _encode_text_prompt(self, prompt, device, num_images_per_prompt, d...
method _encode_image_prompt (line 278) | def _encode_image_prompt(self, prompt, device, num_images_per_prompt, ...
method decode_latents (line 334) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 348) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 365) | def check_inputs(self, prompt, image, height, width, callback_steps):
method prepare_latents (line 383) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method set_transformer_params (line 400) | def set_transformer_params(self, mix_ratio: float = 0.5, condition_typ...
method __call__ (line 414) | def __call__(
FILE: 4DoF/diffusers/pipelines/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py
class VersatileDiffusionImageVariationPipeline (line 35) | class VersatileDiffusionImageVariationPipeline(DiffusionPipeline):
method __init__ (line 59) | def __init__(
method enable_sequential_cpu_offload (line 78) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 97) | def _execution_device(self):
method _encode_prompt (line 114) | def _encode_prompt(self, prompt, device, num_images_per_prompt, do_cla...
method decode_latents (line 194) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 208) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 226) | def check_inputs(self, image, height, width, callback_steps):
method prepare_latents (line 249) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 267) | def __call__(
FILE: 4DoF/diffusers/pipelines/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py
class VersatileDiffusionTextToImagePipeline (line 34) | class VersatileDiffusionTextToImagePipeline(DiffusionPipeline):
method __init__ (line 62) | def __init__(
method _swap_unet_attention_blocks (line 86) | def _swap_unet_attention_blocks(self):
method remove_unused_weights (line 99) | def remove_unused_weights(self):
method enable_sequential_cpu_offload (line 102) | def enable_sequential_cpu_offload(self, gpu_id=0):
method _execution_device (line 121) | def _execution_device(self):
method _encode_prompt (line 138) | def _encode_prompt(self, prompt, device, num_images_per_prompt, do_cla...
method decode_latents (line 251) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 265) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 283) | def check_inputs(
method prepare_latents (line 331) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method __call__ (line 349) | def __call__(
FILE: 4DoF/diffusers/pipelines/vq_diffusion/pipeline_vq_diffusion.py
class LearnedClassifierFreeSamplingEmbeddings (line 30) | class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin):
method __init__ (line 36) | def __init__(self, learnable: bool, hidden_size: Optional[int] = None,...
class VQDiffusionPipeline (line 52) | class VQDiffusionPipeline(DiffusionPipeline):
method __init__ (line 83) | def __init__(
method _encode_prompt (line 103) | def _encode_prompt(self, prompt, num_images_per_prompt, do_classifier_...
method __call__ (line 167) | def __call__(
method truncate (line 310) | def truncate(self, log_p_x_0: torch.FloatTensor, truncation_rate: floa...
FILE: 4DoF/diffusers/schedulers/scheduling_consistency_models.py
class CMStochasticIterativeSchedulerOutput (line 30) | class CMStochasticIterativeSchedulerOutput(BaseOutput):
class CMStochasticIterativeScheduler (line 43) | class CMStochasticIterativeScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 82) | def __init__(
method index_for_timestep (line 106) | def index_for_timestep(self, timestep, schedule_timesteps=None):
method scale_model_input (line 113) | def scale_model_input(
method sigma_to_t (line 136) | def sigma_to_t(self, sigmas: Union[float, np.ndarray]):
method set_timesteps (line 152) | def set_timesteps(
method _convert_to_karras (line 223) | def _convert_to_karras(self, ramp):
method get_scalings (line 235) | def get_scalings(self, sigma):
method get_scalings_for_boundary_condition (line 242) | def get_scalings_for_boundary_condition(self, sigma):
method step (line 264) | def step(
method add_noise (line 354) | def add_noise(
method __len__ (line 379) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddim.py
class DDIMSchedulerOutput (line 32) | class DDIMSchedulerOutput(BaseOutput):
function betas_for_alpha_bar (line 50) | def betas_for_alpha_bar(
function rescale_zero_terminal_snr (line 94) | def rescale_zero_terminal_snr(betas):
class DDIMScheduler (line 130) | class DDIMScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 190) | def __init__(
method scale_model_input (line 243) | def scale_model_input(self, sample: torch.FloatTensor, timestep: Optio...
method _get_variance (line 257) | def _get_variance(self, timestep, prev_timestep):
method _threshold_sample (line 268) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method set_timesteps (line 302) | def set_timesteps(self, num_inference_steps: int, device: Union[str, t...
method step (line 347) | def step(
method add_noise (line 470) | def add_noise(
method get_velocity (line 494) | def get_velocity(
method __len__ (line 514) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddim_flax.py
class DDIMSchedulerState (line 36) | class DDIMSchedulerState:
method create (line 46) | def create(
class FlaxDDIMSchedulerOutput (line 62) | class FlaxDDIMSchedulerOutput(FlaxSchedulerOutput):
class FlaxDDIMScheduler (line 66) | class FlaxDDIMScheduler(FlaxSchedulerMixin, ConfigMixin):
method has_state (line 109) | def has_state(self):
method __init__ (line 113) | def __init__(
method create_state (line 127) | def create_state(self, common: Optional[CommonSchedulerState] = None) ...
method scale_model_input (line 151) | def scale_model_input(
method set_timesteps (line 165) | def set_timesteps(
method _get_variance (line 187) | def _get_variance(self, state: DDIMSchedulerState, timestep, prev_time...
method step (line 199) | def step(
method add_noise (line 286) | def add_noise(
method get_velocity (line 295) | def get_velocity(
method __len__ (line 304) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddim_inverse.py
class DDIMSchedulerOutput (line 31) | class DDIMSchedulerOutput(BaseOutput):
function betas_for_alpha_bar (line 49) | def betas_for_alpha_bar(
class DDIMInverseScheduler (line 93) | class DDIMInverseScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 134) | def __init__(
method scale_model_input (line 187) | def scale_model_input(self, sample: torch.FloatTensor, timestep: Optio...
method set_timesteps (line 201) | def set_timesteps(self, num_inference_steps: int, device: Union[str, t...
method step (line 225) | def step(
method __len__ (line 282) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddim_parallel.py
class DDIMParallelSchedulerOutput (line 32) | class DDIMParallelSchedulerOutput(BaseOutput):
function betas_for_alpha_bar (line 50) | def betas_for_alpha_bar(
function rescale_zero_terminal_snr (line 95) | def rescale_zero_terminal_snr(betas):
class DDIMParallelScheduler (line 131) | class DDIMParallelScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 193) | def __init__(
method scale_model_input (line 247) | def scale_model_input(self, sample: torch.FloatTensor, timestep: Optio...
method _get_variance (line 261) | def _get_variance(self, timestep, prev_timestep=None):
method _batch_get_variance (line 274) | def _batch_get_variance(self, t, prev_t):
method _threshold_sample (line 286) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method set_timesteps (line 321) | def set_timesteps(self, num_inference_steps: int, device: Union[str, t...
method step (line 366) | def step(
method batch_step_no_noise (line 488) | def batch_step_no_noise(
method add_noise (line 597) | def add_noise(
method get_velocity (line 621) | def get_velocity(
method __len__ (line 641) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddpm.py
class DDPMSchedulerOutput (line 30) | class DDPMSchedulerOutput(BaseOutput):
function betas_for_alpha_bar (line 47) | def betas_for_alpha_bar(
class DDPMScheduler (line 91) | class DDPMScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 145) | def __init__(
method scale_model_input (line 195) | def scale_model_input(self, sample: torch.FloatTensor, timestep: Optio...
method set_timesteps (line 209) | def set_timesteps(
method _get_variance (line 284) | def _get_variance(self, t, predicted_variance=None, variance_type=None):
method _threshold_sample (line 324) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method step (line 358) | def step(
method add_noise (line 454) | def add_noise(
method get_velocity (line 477) | def get_velocity(
method __len__ (line 497) | def __len__(self):
method previous_timestep (line 500) | def previous_timestep(self, timestep):
FILE: 4DoF/diffusers/schedulers/scheduling_ddpm_flax.py
class DDPMSchedulerState (line 36) | class DDPMSchedulerState:
method create (line 45) | def create(cls, common: CommonSchedulerState, init_noise_sigma: jnp.nd...
class FlaxDDPMSchedulerOutput (line 50) | class FlaxDDPMSchedulerOutput(FlaxSchedulerOutput):
class FlaxDDPMScheduler (line 54) | class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin):
method has_state (line 92) | def has_state(self):
method __init__ (line 96) | def __init__(
method create_state (line 110) | def create_state(self, common: Optional[CommonSchedulerState] = None) ...
method scale_model_input (line 125) | def scale_model_input(
method set_timesteps (line 139) | def set_timesteps(
method _get_variance (line 162) | def _get_variance(self, state: DDPMSchedulerState, t, predicted_varian...
method step (line 195) | def step(
method add_noise (line 280) | def add_noise(
method get_velocity (line 289) | def get_velocity(
method __len__ (line 298) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_ddpm_parallel.py
class DDPMParallelSchedulerOutput (line 31) | class DDPMParallelSchedulerOutput(BaseOutput):
function betas_for_alpha_bar (line 49) | def betas_for_alpha_bar(
class DDPMParallelScheduler (line 93) | class DDPMParallelScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 149) | def __init__(
method scale_model_input (line 200) | def scale_model_input(self, sample: torch.FloatTensor, timestep: Optio...
method set_timesteps (line 215) | def set_timesteps(
method _get_variance (line 291) | def _get_variance(self, t, predicted_variance=None, variance_type=None):
method _threshold_sample (line 332) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method step (line 366) | def step(
method batch_step_no_noise (line 462) | def batch_step_no_noise(
method add_noise (line 543) | def add_noise(
method get_velocity (line 567) | def get_velocity(
method __len__ (line 587) | def __len__(self):
method previous_timestep (line 591) | def previous_timestep(self, timestep):
FILE: 4DoF/diffusers/schedulers/scheduling_deis_multistep.py
function betas_for_alpha_bar (line 29) | def betas_for_alpha_bar(
class DEISMultistepScheduler (line 73) | class DEISMultistepScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 138) | def __init__(
method set_timesteps (line 202) | def set_timesteps(self, num_inference_steps: int, device: Union[str, t...
method _threshold_sample (line 261) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method convert_model_output (line 295) | def convert_model_output(
method deis_first_order_update (line 333) | def deis_first_order_update(
method multistep_deis_second_order_update (line 363) | def multistep_deis_second_order_update(
method multistep_deis_third_order_update (line 405) | def multistep_deis_third_order_update(
method step (line 462) | def step(
method scale_model_input (line 530) | def scale_model_input(self, sample: torch.FloatTensor, *args, **kwargs...
method add_noise (line 544) | def add_noise(
method __len__ (line 567) | def __len__(self):
FILE: 4DoF/diffusers/schedulers/scheduling_dpmsolver_multistep.py
function betas_for_alpha_bar (line 29) | def betas_for_alpha_bar(
class DPMSolverMultistepScheduler (line 73) | class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
method __init__ (line 165) | def __init__(
method set_timesteps (line 231) | def set_timesteps(self, num_inference_steps: int = None, device: Union...
method _threshold_sample (line 292) | def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatT...
method _sigma_to_t (line 327) | d
Copy disabled (too large)
Download .json
Condensed preview — 768 files, each showing path, character count, and a content snippet. Download the .json file for the full structured content (10,631K chars).
[
{
"path": ".gitignore",
"chars": 2000,
"preview": ".idea/\n\ntraining/\nlightning_logs/\nimage_log/\n\n*.pth\n*.pt\n*.ckpt\n*.safetensors\n\ngradio_pose2image_private.py\ngradio_canny"
},
{
"path": "3drecon/configs/neus_36.yaml",
"chars": 478,
"preview": "model:\n base_lr: 5.0e-4\n target: renderer.renderer.RendererTrainer\n params:\n total_steps: 2000\n warm_up_steps: "
},
{
"path": "3drecon/ours_GSO_T1/NeuS/grandmother/tensorboard_logs/version_0/hparams.yaml",
"chars": 3,
"preview": "{}\n"
},
{
"path": "3drecon/ours_GSO_T1/NeuS/lion/tensorboard_logs/version_0/hparams.yaml",
"chars": 3,
"preview": "{}\n"
},
{
"path": "3drecon/ours_GSO_T1/NeuS/train/tensorboard_logs/version_0/hparams.yaml",
"chars": 3,
"preview": "{}\n"
},
{
"path": "3drecon/raymarching/__init__.py",
"chars": 26,
"preview": "from .raymarching import *"
},
{
"path": "3drecon/raymarching/backend.py",
"chars": 1361,
"preview": "import os\nfrom torch.utils.cpp_extension import load\n\n_src_path = os.path.dirname(os.path.abspath(__file__))\n\nnvcc_flags"
},
{
"path": "3drecon/raymarching/raymarching.py",
"chars": 14712,
"preview": "import numpy as np\nimport time\n\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Function\nfrom torch.cuda.a"
},
{
"path": "3drecon/raymarching/setup.py",
"chars": 2062,
"preview": "import os\nfrom setuptools import setup\nfrom torch.utils.cpp_extension import BuildExtension, CUDAExtension\n\n_src_path = "
},
{
"path": "3drecon/raymarching/src/bindings.cpp",
"chars": 903,
"preview": "#include <torch/extension.h>\n\n#include \"raymarching.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n // utils\n m.de"
},
{
"path": "3drecon/raymarching/src/raymarching.cu",
"chars": 31326,
"preview": "#include <cuda.h>\n#include <cuda_fp16.h>\n#include <cuda_runtime.h>\n\n#include <ATen/cuda/CUDAContext.h>\n#include <torch/t"
},
{
"path": "3drecon/raymarching/src/raymarching.h",
"chars": 2283,
"preview": "#pragma once\n\n#include <stdint.h>\n#include <torch/torch.h>\n\n\nvoid near_far_from_aabb(const at::Tensor rays_o, const at::"
},
{
"path": "3drecon/renderer/agg_net.py",
"chars": 3741,
"preview": "import torch.nn.functional as F\nimport torch.nn as nn\nimport torch\n\ndef weights_init(m):\n if isinstance(m, nn.Linear)"
},
{
"path": "3drecon/renderer/cost_reg_net.py",
"chars": 3565,
"preview": "import torch.nn as nn\n\nclass ConvBnReLU3D(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size=3, s"
},
{
"path": "3drecon/renderer/dummy_dataset.py",
"chars": 1048,
"preview": "import pytorch_lightning as pl\nfrom torch.utils.data import Dataset\nimport webdataset as wds\nfrom torch.utils.data.distr"
},
{
"path": "3drecon/renderer/feature_net.py",
"chars": 1766,
"preview": "import torch.nn as nn\nimport torch.nn.functional as F\n\nclass ConvBnReLU(nn.Module):\n def __init__(self, in_channels, "
},
{
"path": "3drecon/renderer/neus_networks.py",
"chars": 17899,
"preview": "import math\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport tinycudann as "
},
{
"path": "3drecon/renderer/ngp_renderer.py",
"chars": 29447,
"preview": "import math\nimport trimesh\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom p"
},
{
"path": "3drecon/renderer/renderer.py",
"chars": 28825,
"preview": "import abc\nimport os\nfrom pathlib import Path\n\nimport cv2\nimport numpy as np\nimport pytorch_lightning as pl\nimport torch"
},
{
"path": "3drecon/run_NeuS.py",
"chars": 1141,
"preview": "import os\nimport numpy as np\nfrom tqdm import tqdm\n\n# ours + NeuS\nDATA_DIR = \"/home/xin/data/EscherNet/Data/GSO30\" # G"
},
{
"path": "3drecon/train_renderer.py",
"chars": 7808,
"preview": "import argparse\n\nimport imageio\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom pathlib import Path"
},
{
"path": "3drecon/util.py",
"chars": 1941,
"preview": "import importlib\nimport pickle\nimport numpy as np\nimport cv2\n\ndef instantiate_from_config(config):\n if not \"target\" i"
},
{
"path": "4DoF/CN_encoder.py",
"chars": 1198,
"preview": "from transformers import ConvNextV2Model\nimport torch\nfrom typing import Optional\nimport einops\n\nclass CN_encoder(ConvNe"
},
{
"path": "4DoF/dataset.py",
"chars": 10218,
"preview": "import os\nimport math\nfrom pathlib import Path\nimport torch\nimport torchvision\nfrom torch.utils.data import Dataset, Dat"
},
{
"path": "4DoF/diffusers/__init__.py",
"chars": 9329,
"preview": "__version__ = \"0.18.2\"\n\nfrom .configuration_utils import ConfigMixin\nfrom .utils import (\n OptionalDependencyNotAvail"
},
{
"path": "4DoF/diffusers/commands/__init__.py",
"chars": 920,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/commands/diffusers_cli.py",
"chars": 1200,
"preview": "#!/usr/bin/env python\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License,"
},
{
"path": "4DoF/diffusers/commands/env.py",
"chars": 2870,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/configuration_utils.py",
"chars": 30369,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "4DoF/diffusers/dependency_versions_check.py",
"chars": 1756,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/dependency_versions_table.py",
"chars": 1446,
"preview": "# THIS FILE HAS BEEN AUTOGENERATED. To update:\n# 1. modify the `_deps` dict in setup.py\n# 2. run `make deps_table_update"
},
{
"path": "4DoF/diffusers/experimental/__init__.py",
"chars": 38,
"preview": "from .rl import ValueGuidedRLPipeline\n"
},
{
"path": "4DoF/diffusers/experimental/rl/__init__.py",
"chars": 57,
"preview": "from .value_guided_sampling import ValueGuidedRLPipeline\n"
},
{
"path": "4DoF/diffusers/experimental/rl/value_guided_sampling.py",
"chars": 6133,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/image_processor.py",
"chars": 13700,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/loaders.py",
"chars": 75147,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/__init__.py",
"chars": 1446,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/activations.py",
"chars": 297,
"preview": "from torch import nn\n\n\ndef get_activation(act_fn):\n if act_fn in [\"swish\", \"silu\"]:\n return nn.SiLU()\n elif"
},
{
"path": "4DoF/diffusers/models/attention.py",
"chars": 15347,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/attention_flax.py",
"chars": 18142,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/attention_processor.py",
"chars": 72632,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/autoencoder_kl.py",
"chars": 17913,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/controlnet.py",
"chars": 33076,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/controlnet_flax.py",
"chars": 16466,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/cross_attention.py",
"chars": 4875,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/dual_transformer_2d.py",
"chars": 7297,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/embeddings.py",
"chars": 19390,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/embeddings_flax.py",
"chars": 3443,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/modeling_flax_pytorch_utils.py",
"chars": 4601,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lice"
},
{
"path": "4DoF/diffusers/models/modeling_flax_utils.py",
"chars": 26071,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lice"
},
{
"path": "4DoF/diffusers/models/modeling_pytorch_flax_utils.py",
"chars": 6974,
"preview": "# coding=utf-8\r\n# Copyright 2023 The HuggingFace Inc. team.\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"L"
},
{
"path": "4DoF/diffusers/models/modeling_utils.py",
"chars": 46581,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "4DoF/diffusers/models/prior_transformer.py",
"chars": 16574,
"preview": "from dataclasses import dataclass\nfrom typing import Dict, Optional, Union\n\nimport torch\nimport torch.nn.functional as F"
},
{
"path": "4DoF/diffusers/models/resnet.py",
"chars": 35294,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n# `TemporalConvLayer` Copyright 2023 Alibaba DAMO-VILAB, The"
},
{
"path": "4DoF/diffusers/models/resnet_flax.py",
"chars": 4021,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/t5_film_transformer.py",
"chars": 11824,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/transformer_2d.py",
"chars": 17532,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/transformer_temporal.py",
"chars": 7539,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_1d.py",
"chars": 10790,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_1d_blocks.py",
"chars": 24487,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_2d.py",
"chars": 15436,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_2d_blocks.py",
"chars": 119833,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_2d_blocks_flax.py",
"chars": 14206,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_2d_condition.py",
"chars": 50746,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_2d_condition_flax.py",
"chars": 16160,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_3d_blocks.py",
"chars": 23974,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/unet_3d_condition.py",
"chars": 28251,
"preview": "# Copyright 2023 Alibaba DAMO-VILAB and The HuggingFace Team. All rights reserved.\n# Copyright 2023 The ModelScope Team."
},
{
"path": "4DoF/diffusers/models/vae.py",
"chars": 15930,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/vae_flax.py",
"chars": 31945,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/models/vq_model.py",
"chars": 6852,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/optimization.py",
"chars": 14546,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lice"
},
{
"path": "4DoF/diffusers/pipeline_utils.py",
"chars": 1147,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/__init__.py",
"chars": 7009,
"preview": "from ..utils import (\n OptionalDependencyNotAvailable,\n is_flax_available,\n is_invisible_watermark_available,\n "
},
{
"path": "4DoF/diffusers/pipelines/alt_diffusion/__init__.py",
"chars": 1346,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport PIL\nfrom PIL impor"
},
{
"path": "4DoF/diffusers/pipelines/alt_diffusion/modeling_roberta_series.py",
"chars": 5580,
"preview": "from dataclasses import dataclass\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import nn\nfrom transformer"
},
{
"path": "4DoF/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py",
"chars": 40561,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py",
"chars": 41786,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/audio_diffusion/__init__.py",
"chars": 82,
"preview": "from .mel import Mel\nfrom .pipeline_audio_diffusion import AudioDiffusionPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/audio_diffusion/mel.py",
"chars": 5381,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/audio_diffusion/pipeline_audio_diffusion.py",
"chars": 10699,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/audioldm/__init__.py",
"chars": 503,
"preview": "from ...utils import (\n OptionalDependencyNotAvailable,\n is_torch_available,\n is_transformers_available,\n is"
},
{
"path": "4DoF/diffusers/pipelines/audioldm/pipeline_audioldm.py",
"chars": 28856,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/consistency_models/__init__.py",
"chars": 66,
"preview": "from .pipeline_consistency_models import ConsistencyModelPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/consistency_models/pipeline_consistency_models.py",
"chars": 15805,
"preview": "from typing import Callable, List, Optional, Union\n\nimport torch\n\nfrom ...models import UNet2DModel\nfrom ...schedulers i"
},
{
"path": "4DoF/diffusers/pipelines/controlnet/__init__.py",
"chars": 817,
"preview": "from ...utils import (\n OptionalDependencyNotAvailable,\n is_flax_available,\n is_torch_available,\n is_transfo"
},
{
"path": "4DoF/diffusers/pipelines/controlnet/multicontrolnet.py",
"chars": 9229,
"preview": "import os\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom torch import nn\n\nfrom "
},
{
"path": "4DoF/diffusers/pipelines/controlnet/pipeline_controlnet.py",
"chars": 52389,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py",
"chars": 57178,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py",
"chars": 68724,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py",
"chars": 23352,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/dance_diffusion/__init__.py",
"chars": 61,
"preview": "from .pipeline_dance_diffusion import DanceDiffusionPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py",
"chars": 5637,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/ddim/__init__.py",
"chars": 40,
"preview": "from .pipeline_ddim import DDIMPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/ddim/pipeline_ddim.py",
"chars": 5582,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/ddpm/__init__.py",
"chars": 40,
"preview": "from .pipeline_ddpm import DDPMPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/ddpm/pipeline_ddpm.py",
"chars": 4495,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/__init__.py",
"chars": 2163,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport PIL\n\nfrom ...utils"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if.py",
"chars": 39510,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py",
"chars": 44112,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py",
"chars": 49281,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py",
"chars": 49324,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py",
"chars": 54304,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py",
"chars": 44094,
"preview": "import html\nimport inspect\nimport re\nimport urllib.parse as ul\nfrom typing import Any, Callable, Dict, List, Optional, U"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/safety_checker.py",
"chars": 2117,
"preview": "import numpy as np\nimport torch\nimport torch.nn as nn\nfrom transformers import CLIPConfig, CLIPVisionModelWithProjection"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/timesteps.py",
"chars": 5164,
"preview": "fast27_timesteps = [\n 999,\n 800,\n 799,\n 600,\n 599,\n 500,\n 400,\n 399,\n 377,\n 355,\n 333,\n"
},
{
"path": "4DoF/diffusers/pipelines/deepfloyd_if/watermark.py",
"chars": 1595,
"preview": "from typing import List\n\nimport PIL\nimport torch\nfrom PIL import Image\n\nfrom ...configuration_utils import ConfigMixin\nf"
},
{
"path": "4DoF/diffusers/pipelines/dit/__init__.py",
"chars": 38,
"preview": "from .pipeline_dit import DiTPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/dit/pipeline_dit.py",
"chars": 8478,
"preview": "# Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)\n# William Peebles and Saining Xie\n#\n# Copyright (c) 2021 Op"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/__init__.py",
"chars": 749,
"preview": "from ...utils import (\n OptionalDependencyNotAvailable,\n is_torch_available,\n is_transformers_available,\n is"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky.py",
"chars": 20118,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py",
"chars": 23810,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py",
"chars": 29827,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py",
"chars": 25227,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky/text_encoder.py",
"chars": 1022,
"preview": "import torch\nfrom transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel\n\n\nclass MCLIPConfig(XLMRobertaC"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/__init__.py",
"chars": 519,
"preview": "from .pipeline_kandinsky2_2 import KandinskyV22Pipeline\nfrom .pipeline_kandinsky2_2_controlnet import KandinskyV22Contro"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py",
"chars": 13864,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py",
"chars": 16649,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py",
"chars": 19899,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py",
"chars": 17907,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py",
"chars": 23465,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py",
"chars": 24706,
"preview": "from typing import List, Optional, Union\n\nimport PIL\nimport torch\nfrom transformers import CLIPImageProcessor, CLIPTextM"
},
{
"path": "4DoF/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py",
"chars": 27148,
"preview": "from typing import List, Optional, Union\n\nimport PIL\nimport torch\nfrom transformers import CLIPImageProcessor, CLIPTextM"
},
{
"path": "4DoF/diffusers/pipelines/latent_diffusion/__init__.py",
"chars": 243,
"preview": "from ...utils import is_transformers_available\nfrom .pipeline_latent_diffusion_superresolution import LDMSuperResolution"
},
{
"path": "4DoF/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py",
"chars": 32468,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py",
"chars": 6959,
"preview": "import inspect\nfrom typing import List, Optional, Tuple, Union\n\nimport numpy as np\nimport PIL\nimport torch\nimport torch."
},
{
"path": "4DoF/diffusers/pipelines/latent_diffusion_uncond/__init__.py",
"chars": 58,
"preview": "from .pipeline_latent_diffusion_uncond import LDMPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py",
"chars": 4888,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/onnx_utils.py",
"chars": 8282,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "4DoF/diffusers/pipelines/paint_by_example/__init__.py",
"chars": 379,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport PIL\nfrom PIL impor"
},
{
"path": "4DoF/diffusers/pipelines/paint_by_example/image_encoder.py",
"chars": 2441,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py",
"chars": 30649,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/pipeline_flax_utils.py",
"chars": 25924,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "4DoF/diffusers/pipelines/pipeline_utils.py",
"chars": 72441,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "4DoF/diffusers/pipelines/pndm/__init__.py",
"chars": 40,
"preview": "from .pipeline_pndm import PNDMPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/pndm/pipeline_pndm.py",
"chars": 4178,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/repaint/__init__.py",
"chars": 46,
"preview": "from .pipeline_repaint import RePaintPipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/repaint/pipeline_repaint.py",
"chars": 7793,
"preview": "# Copyright 2023 ETH Zurich Computer Vision Lab and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Ap"
},
{
"path": "4DoF/diffusers/pipelines/score_sde_ve/__init__.py",
"chars": 54,
"preview": "from .pipeline_score_sde_ve import ScoreSdeVePipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/score_sde_ve/pipeline_score_sde_ve.py",
"chars": 4353,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/semantic_stable_diffusion/__init__.py",
"chars": 1144,
"preview": "from dataclasses import dataclass\nfrom enum import Enum\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimp"
},
{
"path": "4DoF/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py",
"chars": 39186,
"preview": "import inspect\nimport warnings\nfrom itertools import repeat\nfrom typing import Callable, List, Optional, Union\n\nimport t"
},
{
"path": "4DoF/diffusers/pipelines/shap_e/__init__.py",
"chars": 788,
"preview": "from ...utils import (\n OptionalDependencyNotAvailable,\n is_torch_available,\n is_transformers_available,\n is"
},
{
"path": "4DoF/diffusers/pipelines/shap_e/camera.py",
"chars": 4942,
"preview": "# Copyright 2023 Open AI and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2"
},
{
"path": "4DoF/diffusers/pipelines/shap_e/pipeline_shap_e.py",
"chars": 15962,
"preview": "# Copyright 2023 Open AI and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2"
},
{
"path": "4DoF/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py",
"chars": 13964,
"preview": "# Copyright 2023 Open AI and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2"
},
{
"path": "4DoF/diffusers/pipelines/shap_e/renderer.py",
"chars": 25102,
"preview": "# Copyright 2023 Open AI and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2"
},
{
"path": "4DoF/diffusers/pipelines/spectrogram_diffusion/__init__.py",
"chars": 987,
"preview": "# flake8: noqa\nfrom ...utils import is_note_seq_available, is_transformers_available, is_torch_available\nfrom ...utils i"
},
{
"path": "4DoF/diffusers/pipelines/spectrogram_diffusion/continous_encoder.py",
"chars": 3098,
"preview": "# Copyright 2022 The Music Spectrogram Diffusion Authors.\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n"
},
{
"path": "4DoF/diffusers/pipelines/spectrogram_diffusion/midi_utils.py",
"chars": 25095,
"preview": "# Copyright 2022 The Music Spectrogram Diffusion Authors.\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n"
},
{
"path": "4DoF/diffusers/pipelines/spectrogram_diffusion/notes_encoder.py",
"chars": 2921,
"preview": "# Copyright 2022 The Music Spectrogram Diffusion Authors.\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n"
},
{
"path": "4DoF/diffusers/pipelines/spectrogram_diffusion/pipeline_spectrogram_diffusion.py",
"chars": 8657,
"preview": "# Copyright 2022 The Music Spectrogram Diffusion Authors.\n# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/__init__.py",
"chars": 6164,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport PIL\nfrom PIL impor"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py",
"chars": 69501,
"preview": "# coding=utf-8\n# Copyright 2023 The HuggingFace Inc. team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lice"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py",
"chars": 43535,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py",
"chars": 20922,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_controlnet.py",
"chars": 1257,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py",
"chars": 22777,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py",
"chars": 26220,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py",
"chars": 24226,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py",
"chars": 28398,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py",
"chars": 29009,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint_legacy.py",
"chars": 27673,
"preview": "import inspect\nfrom typing import Callable, List, Optional, Union\n\nimport numpy as np\nimport PIL\nimport torch\nfrom trans"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py",
"chars": 17646,
"preview": "from logging import getLogger\nfrom typing import Any, Callable, List, Optional, Union\n\nimport numpy as np\nimport PIL\nimp"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py",
"chars": 39866,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py",
"chars": 50504,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_controlnet.py",
"chars": 1336,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py",
"chars": 38442,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_diffedit.py",
"chars": 80916,
"preview": "# Copyright 2023 DiffEdit Authors and Pix2Pix Zero Authors and The HuggingFace Team. All rights reserved.\n#\n# Licensed u"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py",
"chars": 23134,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py",
"chars": 42146,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py",
"chars": 55137,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py",
"chars": 42145,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py",
"chars": 40916,
"preview": "# Copyright 2023 The InstructPix2Pix Authors and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apach"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_k_diffusion.py",
"chars": 31678,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py",
"chars": 25330,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py",
"chars": 36026,
"preview": "# Copyright 2023 The Intel Labs Team Authors and the HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apach"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_model_editing.py",
"chars": 40517,
"preview": "# Copyright 2023 TIME Authors and The HuggingFace Team. All rights reserved.\"\n# Licensed under the Apache License, Versi"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_panorama.py",
"chars": 37627,
"preview": "# Copyright 2023 MultiDiffusion Authors and The HuggingFace Team. All rights reserved.\"\n# Licensed under the Apache Lice"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_paradigms.py",
"chars": 43299,
"preview": "# Copyright 2023 ParaDiGMS authors and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License,"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py",
"chars": 62700,
"preview": "# Copyright 2023 Pix2Pix Zero Authors and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache Licen"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py",
"chars": 40488,
"preview": "# Copyright 2023 Susung Hong and The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Versi"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py",
"chars": 38607,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py",
"chars": 46251,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py",
"chars": 40302,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/safety_checker.py",
"chars": 5734,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/safety_checker_flax.py",
"chars": 4486,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py",
"chars": 1890,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_safe/__init__.py",
"chars": 2502,
"preview": "from dataclasses import dataclass\nfrom enum import Enum\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimp"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py",
"chars": 37883,
"preview": "import inspect\nimport warnings\nfrom typing import Callable, List, Optional, Union\n\nimport numpy as np\nimport torch\nfrom "
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_safe/safety_checker.py",
"chars": 5049,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_xl/__init__.py",
"chars": 953,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport PIL\n\nfrom ...utils"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py",
"chars": 41566,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py",
"chars": 46164,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/stable_diffusion_xl/watermark.py",
"chars": 1142,
"preview": "import numpy as np\nimport torch\nfrom imwatermark import WatermarkEncoder\n\n\n# Copied from https://github.com/Stability-AI"
},
{
"path": "4DoF/diffusers/pipelines/stochastic_karras_ve/__init__.py",
"chars": 60,
"preview": "from .pipeline_stochastic_karras_ve import KarrasVePipeline\n"
},
{
"path": "4DoF/diffusers/pipelines/stochastic_karras_ve/pipeline_stochastic_karras_ve.py",
"chars": 5669,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/text_to_video_synthesis/__init__.py",
"chars": 1244,
"preview": "from dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport numpy as np\nimport torch\n\nfrom ...uti"
},
{
"path": "4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py",
"chars": 33145,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py",
"chars": 37412,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "4DoF/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py",
"chars": 27969,
"preview": "import copy\nfrom dataclasses import dataclass\nfrom typing import Callable, List, Optional, Union\n\nimport numpy as np\nimp"
}
]
// ... and 568 more files (download for full content)
About this extraction
This page contains the full source code of the kxhit/EscherNet GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 768 files (126.9 MB), approximately 2.6M tokens, and a symbol index with 6818 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.