Skip to content

k-diffusion 0.1.1.post1

Download JSON

  • PyPI https://pypi.org/project/k-diffusion/
    Repository https://github.com/crowsonkb/k-diffusion

  • py.typed

  • Coverage

    %%{init: {"pie": {"textPosition": 0.85}, "theme": "neutral", "themeVariables": {"pieStrokeWidth": "1px"}}}%%
pie title
    "Typed" : 0
    "Any" : 0
    "Untyped" : 1156
    • 0.0% coverage
    • 0.0% coverage (strict)
    • 1156 typable
      • 0 typed
      • 1156 untyped
      • 0 Any

  • Typables

    %%{init: {"pie": {"textPosition": 0.85}, "theme": "neutral", "themeVariables": {"pieStrokeWidth": "1px"}}}%%
pie title
    "functions" : 408
    "classes" : 580
    "other" : 2
    • 84 functions (+0 overloads)
      • 324 parameters
    • 58 classes
      • 167 methods (+0 overloads)
        • 410 parameters
      • 3 properties
    • 12 modules
      • 2 attrs

Modules

Module Coverage Coverage (strict) Typables Ignores
k_diffusion.augmentation 0.0% 0.0% 37 0
k_diffusion.config 0.0% 0.0% 11 0
k_diffusion.evaluation 0.0% 0.0% 47 0
k_diffusion.external 0.0% 0.0% 98 0
k_diffusion.gns 0.0% 0.0% 26 0
k_diffusion.layers 0.0% 0.0% 146 0
k_diffusion.models.axial_rope 0.0% 0.0% 46 0
k_diffusion.models.flags 0.0% 0.0% 17 0
k_diffusion.models.image_transformer_v1 0.0% 0.0% 177 0
k_diffusion.models.image_v1 0.0% 0.0% 86 0
k_diffusion.sampling 0.0% 0.0% 295 0
k_diffusion.utils 0.0% 0.0% 170 0

Incomplete Annotations

k_diffusion.augmentation (37 missing, 0 any)

Symbol Typable Typed Any
meth KarrasAugmentWrapper.init 2 0 0
meth KarrasAugmentWrapper.forward 6 0 0
meth KarrasAugmentWrapper.param_groups 3 0 0
meth KarrasAugmentWrapper.set_skip_stages 2 0 0
meth KarrasAugmentWrapper.set_patch_size 2 0 0
attr KarrasAugmentWrapper.inner_model 1 0 0
func scale2d 3 0 0
func translate2d 3 0 0
func rotate2d 2 0 0
meth KarrasAugmentationPipeline.init 6 0 0
meth KarrasAugmentationPipeline.call 2 0 0
attr KarrasAugmentationPipeline.a_prob 1 0 0
attr KarrasAugmentationPipeline.a_scale 1 0 0
attr KarrasAugmentationPipeline.a_aniso 1 0 0
attr KarrasAugmentationPipeline.a_trans 1 0 0
attr KarrasAugmentationPipeline.disable_all 1 0 0

k_diffusion.config (11 missing, 0 any)

Symbol Typable Typed Any
func make_sample_density 2 0 0
func round_to_power_of_two 3 0 0
func make_model 2 0 0
func make_denoiser_wrapper 2 0 0
func load_config 2 0 0

k_diffusion.evaluation (47 missing, 0 any)

Symbol Typable Typed Any
func kid 4 0 0
func polynomial_kernel 3 0 0
func fid 4 0 0
meth DINOv2FeatureExtractor.init 3 0 0
meth DINOv2FeatureExtractor.available_models 1 0 0
meth DINOv2FeatureExtractor.forward 2 0 0
attr DINOv2FeatureExtractor.model 1 0 0
attr DINOv2FeatureExtractor.normalize 1 0 0
attr DINOv2FeatureExtractor.size 1 0 0
func compute_features 6 0 0
meth InceptionV3FeatureExtractor.init 2 0 0
meth InceptionV3FeatureExtractor.forward 2 0 0
attr InceptionV3FeatureExtractor.model 1 0 0
attr InceptionV3FeatureExtractor.size 1 0 0
func sqrtm_eig 2 0 0
func squared_mmd 4 0 0
meth CLIPFeatureExtractor.init 3 0 0
meth CLIPFeatureExtractor.available_models 1 0 0
meth CLIPFeatureExtractor.forward 2 0 0
attr CLIPFeatureExtractor.model 1 0 0
attr CLIPFeatureExtractor.normalize 1 0 0
attr CLIPFeatureExtractor.size 1 0 0

k_diffusion.external (98 missing, 0 any)

Symbol Typable Typed Any
meth CompVisDenoiser.init 4 0 0
meth CompVisDenoiser.get_eps 3 0 0
meth DiscreteSchedule.init 3 0 0
meth DiscreteSchedule.get_sigmas 2 0 0
meth DiscreteSchedule.sigma_to_t 3 0 0
meth DiscreteSchedule.t_to_sigma 2 0 0
prop DiscreteSchedule.sigma_min 1 0 0
prop DiscreteSchedule.sigma_max 1 0 0
attr DiscreteSchedule.quantize 1 0 0
meth VDenoiser.init 2 0 0
meth VDenoiser.get_scalings 2 0 0
meth VDenoiser.sigma_to_t 2 0 0
meth VDenoiser.t_to_sigma 2 0 0
meth VDenoiser.loss 5 0 0
meth VDenoiser.forward 4 0 0
attr VDenoiser.inner_model 1 0 0
attr VDenoiser.sigma_data 1 0 0
meth CompVisVDenoiser.init 4 0 0
meth CompVisVDenoiser.get_v 5 0 0
meth DiscreteVDDPMDenoiser.init 4 0 0
meth DiscreteVDDPMDenoiser.get_scalings 2 0 0
meth DiscreteVDDPMDenoiser.get_v 3 0 0
meth DiscreteVDDPMDenoiser.loss 5 0 0
meth DiscreteVDDPMDenoiser.forward 4 0 0
attr DiscreteVDDPMDenoiser.inner_model 1 0 0
attr DiscreteVDDPMDenoiser.sigma_data 1 0 0
meth DiscreteEpsDDPMDenoiser.init 4 0 0
meth DiscreteEpsDDPMDenoiser.get_scalings 2 0 0
meth DiscreteEpsDDPMDenoiser.get_eps 3 0 0
meth DiscreteEpsDDPMDenoiser.loss 5 0 0
meth DiscreteEpsDDPMDenoiser.forward 4 0 0
attr DiscreteEpsDDPMDenoiser.inner_model 1 0 0
attr DiscreteEpsDDPMDenoiser.sigma_data 1 0 0
meth OpenAIDenoiser.init 6 0 0
meth OpenAIDenoiser.get_eps 3 0 0
attr OpenAIDenoiser.has_learned_sigmas 1 0 0

k_diffusion.gns (26 missing, 0 any)

Symbol Typable Typed Any
meth DDPGradientStatsHook.init 2 0 0
meth DDPGradientStatsHook._clear_state 1 0 0
meth DDPGradientStatsHook._hook_fn 3 0 0
meth DDPGradientStatsHook.get_stats 1 0 0
meth GradientNoiseScale.init 3 0 0
meth GradientNoiseScale.state_dict 1 0 0
meth GradientNoiseScale.load_state_dict 2 0 0
meth GradientNoiseScale.update 5 0 0
meth GradientNoiseScale.get_gns 1 0 0
meth GradientNoiseScale.get_stats 1 0 0
attr GradientNoiseScale.beta 1 0 0
attr GradientNoiseScale.eps 1 0 0
attr GradientNoiseScale.ema_sq_norm 1 0 0
attr GradientNoiseScale.ema_var 1 0 0
attr GradientNoiseScale.beta_cumprod 1 0 0
attr GradientNoiseScale.gradient_noise_scale 1 0 0

k_diffusion.layers (146 missing, 0 any)

Symbol Typable Typed Any
meth UNet.init 4 0 0
meth UNet.forward 3 0 0
attr UNet.d_blocks 1 0 0
attr UNet.u_blocks 1 0 0
attr UNet.skip_stages 1 0 0
meth DenoiserWithVariance.loss 5 0 0
meth Downsample2d.init 3 0 0
meth Downsample2d.forward 2 0 0
attr Downsample2d.pad_mode 1 0 0
attr Downsample2d.pad 1 0 0
meth ConditionedResidualBlock.init 3 0 0
meth ConditionedResidualBlock.forward 3 0 0
attr ConditionedResidualBlock.main 1 0 0
attr ConditionedResidualBlock.skip 1 0 0
meth ConditionedSequential.forward 3 0 0
meth SelfAttention2d.init 5 0 0
meth SelfAttention2d.forward 3 0 0
attr SelfAttention2d.norm_in 1 0 0
attr SelfAttention2d.n_head 1 0 0
attr SelfAttention2d.qkv_proj 1 0 0
attr SelfAttention2d.out_proj 1 0 0
attr SelfAttention2d.dropout 1 0 0
func dct 2 0 0
meth Denoiser.init 5 0 0
meth Denoiser._weighting_soft_min_snr 2 0 0
meth Denoiser._weighting_snr 2 0 0
meth Denoiser.get_scalings 2 0 0
meth Denoiser.loss 5 0 0
meth Denoiser.forward 4 0 0
attr Denoiser.inner_model 1 0 0
attr Denoiser.sigma_data 1 0 0
attr Denoiser.scales 1 0 0
attr Denoiser.weighting 1 0 0
func freq_weight_1d 5 0 0
meth Upsample2d.init 3 0 0
meth Upsample2d.forward 2 0 0
attr Upsample2d.pad_mode 1 0 0
attr Upsample2d.pad 1 0 0
meth UnconditionedModule.init 2 0 0
meth UnconditionedModule.forward 3 0 0
attr UnconditionedModule.module 1 0 0
meth SimpleLossDenoiser.loss 5 0 0
meth CrossAttention2d.init 8 0 0
meth CrossAttention2d.forward 3 0 0
attr CrossAttention2d.cond_key 1 0 0
attr CrossAttention2d.cond_key_padding 1 0 0
attr CrossAttention2d.norm_enc 1 0 0
attr CrossAttention2d.norm_dec 1 0 0
attr CrossAttention2d.n_head 1 0 0
attr CrossAttention2d.q_proj 1 0 0
attr CrossAttention2d.kv_proj 1 0 0
attr CrossAttention2d.out_proj 1 0 0
attr CrossAttention2d.dropout 1 0 0
meth AdaGN.init 6 0 0
meth AdaGN.forward 3 0 0
attr AdaGN.num_groups 1 0 0
attr AdaGN.eps 1 0 0
attr AdaGN.cond_key 1 0 0
attr AdaGN.mapper 1 0 0
func freq_weight_nd 5 0 0
meth FourierFeatures.init 4 0 0
meth FourierFeatures.forward 2 0 0
meth ResidualBlock.init 3 0 0
meth ResidualBlock.forward 2 0 0
attr ResidualBlock.main 1 0 0
attr ResidualBlock.skip 1 0 0

k_diffusion.models.axial_rope (46 missing, 0 any)

Symbol Typable Typed Any
meth AxialRoPE.init 5 0 0
meth AxialRoPE.extra_repr 1 0 0
meth AxialRoPE.get_freqs 2 0 0
meth AxialRoPE.forward 3 0 0
attr AxialRoPE.n_heads 1 0 0
attr AxialRoPE.start_index 1 0 0
attr AxialRoPE.freqs_h 1 0 0
attr AxialRoPE.freqs_w 1 0 0
func apply_rotary_emb 5 0 0
func freqs_pixel_log 2 0 0
func bounding_box 4 0 0
func rotate_half 2 0 0
func freqs_pixel 2 0 0
func make_axial_pos 7 0 0
func make_grid 3 0 0
func centers 6 0 0

k_diffusion.models.flags (17 missing, 0 any)

Symbol Typable Typed Any
meth compile_wrap.init 4 0 0
meth compile_wrap.call 3 0 0
prop compile_wrap.compiled_function 1 0 0
attr compile_wrap.function 1 0 0
attr compile_wrap.args 1 0 0
attr compile_wrap.kwargs 1 0 0
func checkpointing 2 0 0
func get_checkpointing 1 0 0
func get_use_flash_attention_2 1 0 0
func get_use_compile 1 0 0
attr state 1 0 0

k_diffusion.models.image_transformer_v1 (177 missing, 0 any)

Symbol Typable Typed Any
meth MappingFeedForwardBlock.init 4 0 0
meth MappingFeedForwardBlock.forward 2 0 0
attr MappingFeedForwardBlock.norm 1 0 0
attr MappingFeedForwardBlock.up_proj 1 0 0
attr MappingFeedForwardBlock.act 1 0 0
attr MappingFeedForwardBlock.dropout 1 0 0
attr MappingFeedForwardBlock.down_proj 1 0 0
meth RMSNorm.init 3 0 0
meth RMSNorm.extra_repr 1 0 0
meth RMSNorm.forward 2 0 0
attr RMSNorm.eps 1 0 0
attr RMSNorm.scale 1 0 0
func tag_module 3 0 0
func checkpoint_helper 4 0 0
meth TransformerBlock.init 5 0 0
meth TransformerBlock.forward 5 0 0
attr TransformerBlock.self_attn 1 0 0
attr TransformerBlock.ff 1 0 0
meth SelfAttentionBlock.init 4 0 0
meth SelfAttentionBlock.extra_repr 1 0 0
meth SelfAttentionBlock.forward 5 0 0
attr SelfAttentionBlock.d_head 1 0 0
attr SelfAttentionBlock.n_heads 1 0 0
attr SelfAttentionBlock.norm 1 0 0
attr SelfAttentionBlock.qkv_proj 1 0 0
attr SelfAttentionBlock.qk_norm 1 0 0
attr SelfAttentionBlock.pos_emb 1 0 0
attr SelfAttentionBlock.dropout 1 0 0
attr SelfAttentionBlock.out_proj 1 0 0
meth Patching.init 3 0 0
meth Patching.extra_repr 1 0 0
meth Patching.forward 3 0 0
attr Patching.features 1 0 0
attr Patching.patch_size 1 0 0
attr Patching.d_out 1 0 0
meth ImageTransformerDenoiserModelV1.init 10 0 0
meth ImageTransformerDenoiserModelV1.proj_ 1 0 0
meth ImageTransformerDenoiserModelV1.param_groups 3 0 0
meth ImageTransformerDenoiserModelV1.forward 5 0 0
attr ImageTransformerDenoiserModelV1.sigma_data 1 0 0
attr ImageTransformerDenoiserModelV1.num_classes 1 0 0
attr ImageTransformerDenoiserModelV1.patch_in 1 0 0
attr ImageTransformerDenoiserModelV1.patch_out 1 0 0
attr ImageTransformerDenoiserModelV1.time_emb 1 0 0
attr ImageTransformerDenoiserModelV1.time_in_proj 1 0 0
attr ImageTransformerDenoiserModelV1.aug_emb 1 0 0
attr ImageTransformerDenoiserModelV1.aug_in_proj 1 0 0
attr ImageTransformerDenoiserModelV1.class_emb 1 0 0
attr ImageTransformerDenoiserModelV1.mapping 1 0 0
attr ImageTransformerDenoiserModelV1.in_proj 1 0 0
attr ImageTransformerDenoiserModelV1.blocks 1 0 0
attr ImageTransformerDenoiserModelV1.out_norm 1 0 0
attr ImageTransformerDenoiserModelV1.out_proj 1 0 0
meth GEGLU.forward 2 0 0
meth FeedForwardBlock.init 4 0 0
meth FeedForwardBlock.forward 3 0 0
attr FeedForwardBlock.norm 1 0 0
attr FeedForwardBlock.up_proj 1 0 0
attr FeedForwardBlock.act 1 0 0
attr FeedForwardBlock.dropout 1 0 0
attr FeedForwardBlock.down_proj 1 0 0
func zero_init 2 0 0
meth QKNorm.init 4 0 0
meth QKNorm.extra_repr 1 0 0
meth QKNorm.proj_ 1 0 0
meth QKNorm.forward 2 0 0
attr QKNorm.eps 1 0 0
attr QKNorm.max_scale 1 0 0
attr QKNorm.scale 1 0 0
func filter_params 3 0 0
func tag_param 3 0 0
meth MappingNetwork.init 5 0 0
meth MappingNetwork.forward 2 0 0
attr MappingNetwork.in_norm 1 0 0
attr MappingNetwork.blocks 1 0 0
attr MappingNetwork.out_norm 1 0 0
func apply_wd 2 0 0
func geglu 2 0 0
func scaled_dot_product_attention 6 0 0
func rms_norm 4 0 0
meth Unpatching.init 3 0 0
meth Unpatching.extra_repr 1 0 0
meth Unpatching.forward 4 0 0
attr Unpatching.features 1 0 0
attr Unpatching.patch_size 1 0 0
attr Unpatching.d_in 1 0 0
meth AdaRMSNorm.init 4 0 0
meth AdaRMSNorm.extra_repr 1 0 0
meth AdaRMSNorm.forward 3 0 0
attr AdaRMSNorm.eps 1 0 0
attr AdaRMSNorm.linear 1 0 0

k_diffusion.models.image_v1 (86 missing, 0 any)

Symbol Typable Typed Any
meth UBlock.init 13 0 0
meth UBlock.forward 4 0 0
meth UBlock.set_upsample 2 0 0
meth ImageDenoiserModelV1.init 14 0 0
meth ImageDenoiserModelV1.param_groups 2 0 0
meth ImageDenoiserModelV1.forward 8 0 0
meth ImageDenoiserModelV1.set_skip_stages 2 0 0
meth ImageDenoiserModelV1.set_patch_size 2 0 0
attr ImageDenoiserModelV1.c_in 1 0 0
attr ImageDenoiserModelV1.channels 1 0 0
attr ImageDenoiserModelV1.unet_cond_dim 1 0 0
attr ImageDenoiserModelV1.patch_size 1 0 0
attr ImageDenoiserModelV1.has_variance 1 0 0
attr ImageDenoiserModelV1.timestep_embed 1 0 0
attr ImageDenoiserModelV1.mapping 1 0 0
attr ImageDenoiserModelV1.proj_in 1 0 0
attr ImageDenoiserModelV1.proj_out 1 0 0
attr ImageDenoiserModelV1.u_net 1 0 0
attr ImageDenoiserModelV1.mapping_cond 1 0 0
func orthogonal_ 2 0 0
meth DBlock.init 13 0 0
meth DBlock.set_downsample 2 0 0
meth MappingNet.init 4 0 0
meth ResConvBlock.init 7 0 0

k_diffusion.sampling (295 missing, 0 any)

Symbol Typable Typed Any
func get_ancestral_step 4 0 0
meth PIDStepSizeController.init 8 0 0
meth PIDStepSizeController.limiter 2 0 0
meth PIDStepSizeController.propose_step 2 0 0
attr PIDStepSizeController.h 1 0 0
attr PIDStepSizeController.b1 1 0 0
attr PIDStepSizeController.b2 1 0 0
attr PIDStepSizeController.b3 1 0 0
attr PIDStepSizeController.accept_safety 1 0 0
attr PIDStepSizeController.eps 1 0 0
attr PIDStepSizeController.errs 1 0 0
func sample_dpm_2 11 0 0
meth BrownianTreeNoiseSampler.init 6 0 0
meth BrownianTreeNoiseSampler.call 3 0 0
attr BrownianTreeNoiseSampler.transform 1 0 0
attr BrownianTreeNoiseSampler.tree 1 0 0
func sample_dpm_2_ancestral 10 0 0
func sample_dpm_adaptive 20 0 0
func get_sigmas_polyexponential 6 0 0
func log_likelihood 8 0 0
meth BatchedBrownianTree.init 6 0 0
meth BatchedBrownianTree.sort 3 0 0
meth BatchedBrownianTree.call 3 0 0
attr BatchedBrownianTree.batched 1 0 0
attr BatchedBrownianTree.trees 1 0 0
func get_sigmas_exponential 5 0 0
func sample_dpmpp_2m_sde 11 0 0
func append_zero 2 0 0
func sample_dpmpp_sde 11 0 0
func get_sigmas_vp 6 0 0
func default_noise_sampler 2 0 0
func sample_euler_ancestral 10 0 0
func to_d 4 0 0
func sample_lms 8 0 0
func get_sigmas_karras 6 0 0
func sample_euler 11 0 0
func sample_dpmpp_3m_sde 10 0 0
func sample_heun 11 0 0
func linear_multistep_coeff 5 0 0
func sample_dpmpp_2s_ancestral 10 0 0
func sample_dpm_fast 12 0 0
func sample_dpmpp_2m 7 0 0
meth DPMSolver.init 5 0 0
meth DPMSolver.t 2 0 0
meth DPMSolver.sigma 2 0 0
meth DPMSolver.eps 7 0 0
meth DPMSolver.dpm_solver_1_step 5 0 0
meth DPMSolver.dpm_solver_2_step 6 0 0
meth DPMSolver.dpm_solver_3_step 7 0 0
meth DPMSolver.dpm_solver_fast 8 0 0
meth DPMSolver.dpm_solver_adaptive 15 0 0
attr DPMSolver.model 1 0 0
attr DPMSolver.extra_args 1 0 0
attr DPMSolver.eps_callback 1 0 0
attr DPMSolver.info_callback 1 0 0

k_diffusion.utils (170 missing, 0 any)

Symbol Typable Typed Any
func train_mode 3 0 0
meth FolderOfImages.init 3 0 0
meth FolderOfImages.repr 1 0 0
meth FolderOfImages.len 1 0 0
meth FolderOfImages.getitem 2 0 0
attr FolderOfImages.root 1 0 0
attr FolderOfImages.transform 1 0 0
attr FolderOfImages.paths 1 0 0
func enable_stratified_accelerate 3 0 0
func rand_cosine_interpolated 10 0 0
func get_safetensors_metadata 2 0 0
meth CSVLogger.init 3 0 0
meth CSVLogger.write 2 0 0
attr CSVLogger.filename 1 0 0
attr CSVLogger.columns 1 0 0
attr CSVLogger.file 1 0 0
func tf32_mode 3 0 0
func n_params 2 0 0
func hf_datasets_augs_helper 5 0 0
meth ExponentialLR.init 8 0 0
meth ExponentialLR.get_lr 1 0 0
meth ExponentialLR._get_closed_form_lr 1 0 0
attr ExponentialLR.num_steps 1 0 0
attr ExponentialLR.decay 1 0 0
attr ExponentialLR.warmup 1 0 0
attr ExponentialLR.min_lr 1 0 0
func rand_log_uniform 6 0 0
meth EMAWarmup.init 7 0 0
meth EMAWarmup.state_dict 1 0 0
meth EMAWarmup.load_state_dict 2 0 0
meth EMAWarmup.get_value 1 0 0
meth EMAWarmup.step 1 0 0
attr EMAWarmup.inv_gamma 1 0 0
attr EMAWarmup.power 1 0 0
attr EMAWarmup.min_value 1 0 0
attr EMAWarmup.max_value 1 0 0
attr EMAWarmup.start_at 1 0 0
attr EMAWarmup.last_epoch 1 0 0
meth InverseLR.init 8 0 0
meth InverseLR.get_lr 1 0 0
meth InverseLR._get_closed_form_lr 1 0 0
attr InverseLR.inv_gamma 1 0 0
attr InverseLR.power 1 0 0
attr InverseLR.warmup 1 0 0
attr InverseLR.min_lr 1 0 0
func rand_v_diffusion 7 0 0
meth ConstantLRWithWarmup.init 5 0 0
meth ConstantLRWithWarmup.get_lr 1 0 0
meth ConstantLRWithWarmup._get_closed_form_lr 1 0 0
attr ConstantLRWithWarmup.warmup 1 0 0
func download_file 4 0 0
func from_pil_image 2 0 0
func rand_log_logistic 8 0 0
func append_dims 3 0 0
func rand_split_log_normal 7 0 0
func ema_update 4 0 0
func to_pil_image 2 0 0
func eval_mode 2 0 0
func stratified_uniform 6 0 0
func stratified_with_settings 4 0 0
attr stratified_settings 1 0 0
func rand_log_normal 6 0 0
func enable_stratified 4 0 0
func ema_update_dict 4 0 0

Type-Ignore Comments

No type-ignore comments.