{ "package": "k-diffusion", "version": "0.1.1.post1", "stubs_only": "no", "py_typed": "NO", "pypi": { "upload_time": "2023-12-07T13:27:03.011200Z", "requires_python": null, "size": 30543, "sha256": "04b66b0e95202a60c4c178a482e6d31528b81a1c97011d6023ed3ad559de7166" }, "metadata": { "Metadata-Version": [ "2.1" ], "Name": [ "k-diffusion" ], "Version": [ "0.1.1.post1" ], "Summary": [ "Karras et al. (2022) diffusion models for PyTorch" ], "Home-page": [ "https://github.com/crowsonkb/k-diffusion" ], "Author": [ "Katherine Crowson" ], "Author-email": [ "crowsonkb@gmail.com" ], "License": [ "MIT" ], "Platform": [ "UNKNOWN" ], "Description-Content-Type": [ "text/markdown" ], "Requires-Dist": [ "accelerate", "clean-fid", "clip-anytorch", "dctorch", "einops", "jsonmerge", "kornia", "Pillow", "safetensors", "scikit-image", "scipy", "torch >=2.0", "torchdiffeq", "torchsde", "torchvision", "tqdm", "wandb" ] }, "module_reports": [ { "path": "k_diffusion/augmentation.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.augmentation.KarrasAugmentWrapper", "methods": [ { "kind": "function", "name": "KarrasAugmentWrapper.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "KarrasAugmentWrapper.forward", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "KarrasAugmentWrapper.param_groups", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "KarrasAugmentWrapper.set_skip_stages", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "KarrasAugmentWrapper.set_patch_size", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "KarrasAugmentWrapper.inner_model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 16, "n_typed": 0, "n_any": 0, "n_untyped": 16, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 10, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.augmentation.scale2d", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.augmentation.translate2d", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.augmentation.rotate2d", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "k_diffusion.augmentation.KarrasAugmentationPipeline", "methods": [ { "kind": "function", "name": "KarrasAugmentationPipeline.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "KarrasAugmentationPipeline.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "KarrasAugmentationPipeline.a_prob", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "KarrasAugmentationPipeline.a_scale", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "KarrasAugmentationPipeline.a_aniso", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "KarrasAugmentationPipeline.a_trans", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "KarrasAugmentationPipeline.disable_all", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 5, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.augmentation", "names": [ "k_diffusion.augmentation.KarrasAugmentWrapper", "k_diffusion.augmentation.KarrasAugmentationPipeline", "k_diffusion.augmentation.rotate2d", "k_diffusion.augmentation.scale2d", "k_diffusion.augmentation.translate2d" ], "n_typable": 37, "n_typed": 0, "n_any": 0, "n_untyped": 37, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 5, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 16, "n_classes": 2, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/config.py", "symbol_reports": [ { "kind": "function", "name": "k_diffusion.config.make_sample_density", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.config.round_to_power_of_two", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.config.make_model", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.config.make_denoiser_wrapper", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.config.load_config", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "k_diffusion.config", "names": [ "k_diffusion.config.load_config", "k_diffusion.config.make_denoiser_wrapper", "k_diffusion.config.make_model", "k_diffusion.config.make_sample_density", "k_diffusion.config.round_to_power_of_two" ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/evaluation.py", "symbol_reports": [ { "kind": "function", "name": "k_diffusion.evaluation.kid", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.evaluation.polynomial_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.evaluation.fid", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "k_diffusion.evaluation.DINOv2FeatureExtractor", "methods": [ { "kind": "function", "name": "DINOv2FeatureExtractor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DINOv2FeatureExtractor.available_models", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DINOv2FeatureExtractor.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DINOv2FeatureExtractor.model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DINOv2FeatureExtractor.normalize", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DINOv2FeatureExtractor.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 9, "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.evaluation.compute_features", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "k_diffusion.evaluation.InceptionV3FeatureExtractor", "methods": [ { "kind": "function", "name": "InceptionV3FeatureExtractor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "InceptionV3FeatureExtractor.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "InceptionV3FeatureExtractor.model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "InceptionV3FeatureExtractor.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 6, "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.evaluation.sqrtm_eig", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.evaluation.squared_mmd", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "k_diffusion.evaluation.CLIPFeatureExtractor", "methods": [ { "kind": "function", "name": "CLIPFeatureExtractor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CLIPFeatureExtractor.available_models", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CLIPFeatureExtractor.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CLIPFeatureExtractor.model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CLIPFeatureExtractor.normalize", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CLIPFeatureExtractor.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 9, "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.evaluation", "names": [ "k_diffusion.evaluation.CLIPFeatureExtractor", "k_diffusion.evaluation.DINOv2FeatureExtractor", "k_diffusion.evaluation.InceptionV3FeatureExtractor", "k_diffusion.evaluation.compute_features", "k_diffusion.evaluation.fid", "k_diffusion.evaluation.kid", "k_diffusion.evaluation.polynomial_kernel", "k_diffusion.evaluation.sqrtm_eig", "k_diffusion.evaluation.squared_mmd" ], "n_typable": 47, "n_typed": 0, "n_any": 0, "n_untyped": 47, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 17, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 8, "n_classes": 3, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/external.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.external.CompVisDenoiser", "methods": [ { "kind": "function", "name": "CompVisDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CompVisDenoiser.get_eps", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.external.DiscreteSchedule", "methods": [ { "kind": "function", "name": "DiscreteSchedule.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DiscreteSchedule.get_sigmas", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DiscreteSchedule.sigma_to_t", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DiscreteSchedule.t_to_sigma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "DiscreteSchedule.sigma_min", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "DiscreteSchedule.sigma_max", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "DiscreteSchedule.quantize", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 6, "n_attrs": 1, "n_properties": 2 }, { "kind": "class", "name": "k_diffusion.external.VDenoiser", "methods": [ { "kind": "function", "name": "VDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VDenoiser.get_scalings", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VDenoiser.sigma_to_t", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VDenoiser.t_to_sigma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VDenoiser.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "VDenoiser.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "VDenoiser.inner_model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "VDenoiser.sigma_data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 19, "n_typed": 0, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 11, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.external.CompVisVDenoiser", "methods": [ { "kind": "function", "name": "CompVisVDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CompVisVDenoiser.get_v", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [], "n_typable": 9, "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.external.DiscreteVDDPMDenoiser", "methods": [ { "kind": "function", "name": "DiscreteVDDPMDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "DiscreteVDDPMDenoiser.get_scalings", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DiscreteVDDPMDenoiser.get_v", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DiscreteVDDPMDenoiser.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "DiscreteVDDPMDenoiser.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DiscreteVDDPMDenoiser.inner_model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DiscreteVDDPMDenoiser.sigma_data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 20, "n_typed": 0, "n_any": 0, "n_untyped": 20, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 13, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.external.DiscreteEpsDDPMDenoiser", "methods": [ { "kind": "function", "name": "DiscreteEpsDDPMDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "DiscreteEpsDDPMDenoiser.get_scalings", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DiscreteEpsDDPMDenoiser.get_eps", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DiscreteEpsDDPMDenoiser.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "DiscreteEpsDDPMDenoiser.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DiscreteEpsDDPMDenoiser.inner_model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DiscreteEpsDDPMDenoiser.sigma_data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 20, "n_typed": 0, "n_any": 0, "n_untyped": 20, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 13, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.external.OpenAIDenoiser", "methods": [ { "kind": "function", "name": "OpenAIDenoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "OpenAIDenoiser.get_eps", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "OpenAIDenoiser.has_learned_sigmas", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 1, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.external", "names": [ "k_diffusion.external.CompVisDenoiser", "k_diffusion.external.CompVisVDenoiser", "k_diffusion.external.DiscreteEpsDDPMDenoiser", "k_diffusion.external.DiscreteSchedule", "k_diffusion.external.DiscreteVDDPMDenoiser", "k_diffusion.external.OpenAIDenoiser", "k_diffusion.external.VDenoiser" ], "n_typable": 98, "n_typed": 0, "n_any": 0, "n_untyped": 98, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 26, "n_method_overloads": 26, "n_method_params": 62, "n_classes": 7, "n_attrs": 8, "n_properties": 2, "n_type_ignores": 0 }, { "path": "k_diffusion/gns.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.gns.DDPGradientStatsHook", "methods": [ { "kind": "function", "name": "DDPGradientStatsHook.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DDPGradientStatsHook._clear_state", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DDPGradientStatsHook._hook_fn", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DDPGradientStatsHook.get_stats", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.gns.GradientNoiseScale", "methods": [ { "kind": "function", "name": "GradientNoiseScale.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "GradientNoiseScale.state_dict", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GradientNoiseScale.load_state_dict", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GradientNoiseScale.update", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "GradientNoiseScale.get_gns", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GradientNoiseScale.get_stats", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "GradientNoiseScale.beta", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "GradientNoiseScale.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "GradientNoiseScale.ema_sq_norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "GradientNoiseScale.ema_var", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "GradientNoiseScale.beta_cumprod", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "GradientNoiseScale.gradient_noise_scale", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 19, "n_typed": 0, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 7, "n_attrs": 6, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.gns", "names": [ "k_diffusion.gns.DDPGradientStatsHook", "k_diffusion.gns.GradientNoiseScale" ], "n_typable": 26, "n_typed": 0, "n_any": 0, "n_untyped": 26, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 10, "n_classes": 2, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/layers.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.layers.UNet", "methods": [ { "kind": "function", "name": "UNet.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "UNet.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "UNet.d_blocks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "UNet.u_blocks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "UNet.skip_stages", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.DenoiserWithVariance", "methods": [ { "kind": "function", "name": "DenoiserWithVariance.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 4, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.Downsample2d", "methods": [ { "kind": "function", "name": "Downsample2d.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Downsample2d.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Downsample2d.pad_mode", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Downsample2d.pad", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.ConditionedModule", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.ConditionedResidualBlock", "methods": [ { "kind": "function", "name": "ConditionedResidualBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ConditionedResidualBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ConditionedResidualBlock.main", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ConditionedResidualBlock.skip", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.ConditionedSequential", "methods": [ { "kind": "function", "name": "ConditionedSequential.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.SelfAttention2d", "methods": [ { "kind": "function", "name": "SelfAttention2d.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "SelfAttention2d.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SelfAttention2d.norm_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttention2d.n_head", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttention2d.qkv_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttention2d.out_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttention2d.dropout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.layers.dct", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "k_diffusion.layers.Denoiser", "methods": [ { "kind": "function", "name": "Denoiser.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Denoiser._weighting_soft_min_snr", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Denoiser._weighting_snr", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Denoiser.get_scalings", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Denoiser.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Denoiser.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Denoiser.inner_model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Denoiser.sigma_data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Denoiser.scales", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Denoiser.weighting", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 24, "n_typed": 0, "n_any": 0, "n_untyped": 24, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 14, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.layers.freq_weight_1d", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "k_diffusion.layers.Upsample2d", "methods": [ { "kind": "function", "name": "Upsample2d.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Upsample2d.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Upsample2d.pad_mode", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Upsample2d.pad", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.UnconditionedModule", "methods": [ { "kind": "function", "name": "UnconditionedModule.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "UnconditionedModule.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "UnconditionedModule.module", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 6, "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.SimpleLossDenoiser", "methods": [ { "kind": "function", "name": "SimpleLossDenoiser.loss", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 4, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.CrossAttention2d", "methods": [ { "kind": "function", "name": "CrossAttention2d.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "CrossAttention2d.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CrossAttention2d.cond_key", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.cond_key_padding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.norm_enc", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.norm_dec", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.n_head", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.q_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.kv_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.out_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CrossAttention2d.dropout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 20, "n_typed": 0, "n_any": 0, "n_untyped": 20, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 9, "n_attrs": 9, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.AdaGN", "methods": [ { "kind": "function", "name": "AdaGN.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "AdaGN.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AdaGN.num_groups", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AdaGN.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AdaGN.cond_key", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AdaGN.mapper", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.layers.freq_weight_nd", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "k_diffusion.layers.FourierFeatures", "methods": [ { "kind": "function", "name": "FourierFeatures.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FourierFeatures.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 6, "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.layers.ResidualBlock", "methods": [ { "kind": "function", "name": "ResidualBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ResidualBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ResidualBlock.main", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ResidualBlock.skip", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.layers", "names": [ "k_diffusion.layers.AdaGN", "k_diffusion.layers.ConditionedModule", "k_diffusion.layers.ConditionedResidualBlock", "k_diffusion.layers.ConditionedSequential", "k_diffusion.layers.CrossAttention2d", "k_diffusion.layers.Denoiser", "k_diffusion.layers.DenoiserWithVariance", "k_diffusion.layers.Downsample2d", "k_diffusion.layers.FourierFeatures", "k_diffusion.layers.ResidualBlock", "k_diffusion.layers.SelfAttention2d", "k_diffusion.layers.SimpleLossDenoiser", "k_diffusion.layers.UNet", "k_diffusion.layers.UnconditionedModule", "k_diffusion.layers.Upsample2d", "k_diffusion.layers.dct", "k_diffusion.layers.freq_weight_1d", "k_diffusion.layers.freq_weight_nd" ], "n_typable": 146, "n_typed": 0, "n_any": 0, "n_untyped": 146, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 9, "n_methods": 29, "n_method_overloads": 29, "n_method_params": 71, "n_classes": 15, "n_attrs": 34, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/models/axial_rope.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.models.axial_rope.AxialRoPE", "methods": [ { "kind": "function", "name": "AxialRoPE.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "AxialRoPE.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AxialRoPE.get_freqs", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AxialRoPE.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AxialRoPE.n_heads", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AxialRoPE.start_index", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AxialRoPE.freqs_h", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AxialRoPE.freqs_w", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 15, "n_typed": 0, "n_any": 0, "n_untyped": 15, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 7, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.apply_rotary_emb", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.freqs_pixel_log", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.bounding_box", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.rotate_half", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.freqs_pixel", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.make_axial_pos", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.make_grid", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.models.axial_rope.centers", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "k_diffusion.models.axial_rope", "names": [ "k_diffusion.models.axial_rope.AxialRoPE", "k_diffusion.models.axial_rope.apply_rotary_emb", "k_diffusion.models.axial_rope.bounding_box", "k_diffusion.models.axial_rope.centers", "k_diffusion.models.axial_rope.freqs_pixel", "k_diffusion.models.axial_rope.freqs_pixel_log", "k_diffusion.models.axial_rope.make_axial_pos", "k_diffusion.models.axial_rope.make_grid", "k_diffusion.models.axial_rope.rotate_half" ], "n_typable": 46, "n_typed": 0, "n_any": 0, "n_untyped": 46, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 23, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 7, "n_classes": 1, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/models/flags.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.models.flags.compile_wrap", "methods": [ { "kind": "function", "name": "compile_wrap.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "compile_wrap.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "compile_wrap.compiled_function", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "compile_wrap.function", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "compile_wrap.args", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "compile_wrap.kwargs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 3, "n_properties": 1 }, { "kind": "function", "name": "k_diffusion.models.flags.checkpointing", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.flags.get_checkpointing", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "k_diffusion.models.flags.get_use_flash_attention_2", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "k_diffusion.models.flags.get_use_compile", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "k_diffusion.models.flags.state", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "k_diffusion.models.flags", "names": [ "k_diffusion.models.flags.checkpointing", "k_diffusion.models.flags.compile_wrap", "k_diffusion.models.flags.get_checkpointing", "k_diffusion.models.flags.get_use_compile", "k_diffusion.models.flags.get_use_flash_attention_2", "k_diffusion.models.flags.state" ], "n_typable": 17, "n_typed": 0, "n_any": 0, "n_untyped": 17, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 1, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_classes": 1, "n_attrs": 4, "n_properties": 1, "n_type_ignores": 0 }, { "path": "k_diffusion/models/image_transformer_v1.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.MappingFeedForwardBlock", "methods": [ { "kind": "function", "name": "MappingFeedForwardBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "MappingFeedForwardBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MappingFeedForwardBlock.norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingFeedForwardBlock.up_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingFeedForwardBlock.act", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingFeedForwardBlock.dropout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingFeedForwardBlock.down_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 5, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.RMSNorm", "methods": [ { "kind": "function", "name": "RMSNorm.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "RMSNorm.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "RMSNorm.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RMSNorm.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "RMSNorm.scale", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.tag_module", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.checkpoint_helper", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.TransformerBlock", "methods": [ { "kind": "function", "name": "TransformerBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "TransformerBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TransformerBlock.self_attn", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "TransformerBlock.ff", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 12, "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 8, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.SelfAttentionBlock", "methods": [ { "kind": "function", "name": "SelfAttentionBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "SelfAttentionBlock.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SelfAttentionBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SelfAttentionBlock.d_head", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.n_heads", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.qkv_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.qk_norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.pos_emb", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.dropout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SelfAttentionBlock.out_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 18, "n_typed": 0, "n_any": 0, "n_untyped": 18, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 7, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.Patching", "methods": [ { "kind": "function", "name": "Patching.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Patching.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Patching.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Patching.features", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Patching.patch_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Patching.d_out", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.ImageTransformerDenoiserModelV1", "methods": [ { "kind": "function", "name": "ImageTransformerDenoiserModelV1.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "ImageTransformerDenoiserModelV1.proj_", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ImageTransformerDenoiserModelV1.param_groups", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ImageTransformerDenoiserModelV1.forward", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.sigma_data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.num_classes", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.patch_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.patch_out", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.time_emb", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.time_in_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.aug_emb", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.aug_in_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.class_emb", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.mapping", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.in_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.blocks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.out_norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageTransformerDenoiserModelV1.out_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 33, "n_typed": 0, "n_any": 0, "n_untyped": 33, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 15, "n_attrs": 14, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.GEGLU", "methods": [ { "kind": "function", "name": "GEGLU.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.FeedForwardBlock", "methods": [ { "kind": "function", "name": "FeedForwardBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FeedForwardBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "FeedForwardBlock.norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FeedForwardBlock.up_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FeedForwardBlock.act", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FeedForwardBlock.dropout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FeedForwardBlock.down_proj", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 12, "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.zero_init", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.QKNorm", "methods": [ { "kind": "function", "name": "QKNorm.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "QKNorm.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "QKNorm.proj_", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "QKNorm.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "QKNorm.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "QKNorm.max_scale", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "QKNorm.scale", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.filter_params", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.tag_param", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.MappingNetwork", "methods": [ { "kind": "function", "name": "MappingNetwork.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "MappingNetwork.forward", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MappingNetwork.in_norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingNetwork.blocks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MappingNetwork.out_norm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.apply_wd", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.geglu", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.scaled_dot_product_attention", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.models.image_transformer_v1.rms_norm", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.Unpatching", "methods": [ { "kind": "function", "name": "Unpatching.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Unpatching.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Unpatching.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Unpatching.features", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Unpatching.patch_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Unpatching.d_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_transformer_v1.AdaRMSNorm", "methods": [ { "kind": "function", "name": "AdaRMSNorm.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AdaRMSNorm.extra_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AdaRMSNorm.forward", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AdaRMSNorm.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AdaRMSNorm.linear", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.models.image_transformer_v1", "names": [ "k_diffusion.models.image_transformer_v1.AdaRMSNorm", "k_diffusion.models.image_transformer_v1.FeedForwardBlock", "k_diffusion.models.image_transformer_v1.GEGLU", "k_diffusion.models.image_transformer_v1.ImageTransformerDenoiserModelV1", "k_diffusion.models.image_transformer_v1.MappingFeedForwardBlock", "k_diffusion.models.image_transformer_v1.MappingNetwork", "k_diffusion.models.image_transformer_v1.Patching", "k_diffusion.models.image_transformer_v1.QKNorm", "k_diffusion.models.image_transformer_v1.RMSNorm", "k_diffusion.models.image_transformer_v1.SelfAttentionBlock", "k_diffusion.models.image_transformer_v1.TransformerBlock", "k_diffusion.models.image_transformer_v1.Unpatching", "k_diffusion.models.image_transformer_v1.apply_wd", "k_diffusion.models.image_transformer_v1.checkpoint_helper", "k_diffusion.models.image_transformer_v1.filter_params", "k_diffusion.models.image_transformer_v1.geglu", "k_diffusion.models.image_transformer_v1.rms_norm", "k_diffusion.models.image_transformer_v1.scaled_dot_product_attention", "k_diffusion.models.image_transformer_v1.tag_module", "k_diffusion.models.image_transformer_v1.tag_param", "k_diffusion.models.image_transformer_v1.zero_init" ], "n_typable": 177, "n_typed": 0, "n_any": 0, "n_untyped": 177, "n_functions": 9, "n_function_overloads": 9, "n_function_params": 20, "n_methods": 32, "n_method_overloads": 32, "n_method_params": 66, "n_classes": 12, "n_attrs": 50, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/models/image_v1.py", "symbol_reports": [ { "kind": "class", "name": "k_diffusion.models.image_v1.UBlock", "methods": [ { "kind": "function", "name": "UBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "function", "name": "UBlock.forward", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "UBlock.set_upsample", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 19, "n_typed": 0, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 16, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_v1.ImageDenoiserModelV1", "methods": [ { "kind": "function", "name": "ImageDenoiserModelV1.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 14, "n_overloads": 1, "n_typable": 14, "n_params": 13, "n_function_overloads": 1, "n_function_params": 13 }, { "kind": "function", "name": "ImageDenoiserModelV1.param_groups", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ImageDenoiserModelV1.forward", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "ImageDenoiserModelV1.set_skip_stages", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ImageDenoiserModelV1.set_patch_size", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ImageDenoiserModelV1.c_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.channels", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.unet_cond_dim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.patch_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.has_variance", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.timestep_embed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.mapping", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.proj_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.proj_out", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.u_net", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ImageDenoiserModelV1.mapping_cond", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 39, "n_typed": 0, "n_any": 0, "n_untyped": 39, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 23, "n_attrs": 11, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.models.image_v1.orthogonal_", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "k_diffusion.models.image_v1.DBlock", "methods": [ { "kind": "function", "name": "DBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "function", "name": "DBlock.set_downsample", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 15, "n_typed": 0, "n_any": 0, "n_untyped": 15, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 13, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_v1.MappingNet", "methods": [ { "kind": "function", "name": "MappingNet.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.models.image_v1.ResConvBlock", "methods": [ { "kind": "function", "name": "ResConvBlock.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "properties": [], "attrs": [], "n_typable": 7, "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 6, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.models.image_v1", "names": [ "k_diffusion.models.image_v1.DBlock", "k_diffusion.models.image_v1.ImageDenoiserModelV1", "k_diffusion.models.image_v1.MappingNet", "k_diffusion.models.image_v1.ResConvBlock", "k_diffusion.models.image_v1.UBlock", "k_diffusion.models.image_v1.orthogonal_" ], "n_typable": 86, "n_typed": 0, "n_any": 0, "n_untyped": 86, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 12, "n_method_overloads": 12, "n_method_params": 61, "n_classes": 5, "n_attrs": 11, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/sampling.py", "symbol_reports": [ { "kind": "function", "name": "k_diffusion.sampling.get_ancestral_step", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "k_diffusion.sampling.PIDStepSizeController", "methods": [ { "kind": "function", "name": "PIDStepSizeController.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "PIDStepSizeController.limiter", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PIDStepSizeController.propose_step", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "PIDStepSizeController.h", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.b1", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.b2", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.b3", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.accept_safety", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.eps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "PIDStepSizeController.errs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 19, "n_typed": 0, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 9, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpm_2", "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "class", "name": "k_diffusion.sampling.BrownianTreeNoiseSampler", "methods": [ { "kind": "function", "name": "BrownianTreeNoiseSampler.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "BrownianTreeNoiseSampler.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BrownianTreeNoiseSampler.transform", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BrownianTreeNoiseSampler.tree", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 11, "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpm_2_ancestral", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpm_adaptive", "n_typed": 0, "n_any": 0, "n_untyped": 20, "n_overloads": 1, "n_typable": 20, "n_params": 19, "n_function_overloads": 1, "n_function_params": 19 }, { "kind": "function", "name": "k_diffusion.sampling.get_sigmas_polyexponential", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.sampling.log_likelihood", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "k_diffusion.sampling.BatchedBrownianTree", "methods": [ { "kind": "function", "name": "BatchedBrownianTree.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "BatchedBrownianTree.sort", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BatchedBrownianTree.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BatchedBrownianTree.batched", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BatchedBrownianTree.trees", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 14, "n_typed": 0, "n_any": 0, "n_untyped": 14, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 9, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.sampling.get_sigmas_exponential", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpmpp_2m_sde", "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "k_diffusion.sampling.append_zero", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpmpp_sde", "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "k_diffusion.sampling.get_sigmas_vp", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.sampling.default_noise_sampler", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.sampling.sample_euler_ancestral", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "k_diffusion.sampling.to_d", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.sampling.sample_lms", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "k_diffusion.sampling.get_sigmas_karras", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.sampling.sample_euler", "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpmpp_3m_sde", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "k_diffusion.sampling.sample_heun", "n_typed": 0, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "k_diffusion.sampling.linear_multistep_coeff", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpmpp_2s_ancestral", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpm_fast", "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "k_diffusion.sampling.sample_dpmpp_2m", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "k_diffusion.sampling.DPMSolver", "methods": [ { "kind": "function", "name": "DPMSolver.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "DPMSolver.t", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DPMSolver.sigma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DPMSolver.eps", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "DPMSolver.dpm_solver_1_step", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "DPMSolver.dpm_solver_2_step", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "DPMSolver.dpm_solver_3_step", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "DPMSolver.dpm_solver_fast", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "DPMSolver.dpm_solver_adaptive", "n_typed": 0, "n_any": 0, "n_untyped": 15, "n_overloads": 1, "n_typable": 15, "n_params": 14, "n_function_overloads": 1, "n_function_params": 14 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DPMSolver.model", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DPMSolver.extra_args", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DPMSolver.eps_callback", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DPMSolver.info_callback", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 61, "n_typed": 0, "n_any": 0, "n_untyped": 61, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 48, "n_attrs": 4, "n_properties": 0 } ], "type_ignores": [], "name": "k_diffusion.sampling", "names": [ "k_diffusion.sampling.BatchedBrownianTree", "k_diffusion.sampling.BrownianTreeNoiseSampler", "k_diffusion.sampling.DPMSolver", "k_diffusion.sampling.PIDStepSizeController", "k_diffusion.sampling.append_zero", "k_diffusion.sampling.default_noise_sampler", "k_diffusion.sampling.get_ancestral_step", "k_diffusion.sampling.get_sigmas_exponential", "k_diffusion.sampling.get_sigmas_karras", "k_diffusion.sampling.get_sigmas_polyexponential", "k_diffusion.sampling.get_sigmas_vp", "k_diffusion.sampling.linear_multistep_coeff", "k_diffusion.sampling.log_likelihood", "k_diffusion.sampling.sample_dpm_2", "k_diffusion.sampling.sample_dpm_2_ancestral", "k_diffusion.sampling.sample_dpm_adaptive", "k_diffusion.sampling.sample_dpm_fast", "k_diffusion.sampling.sample_dpmpp_2m", "k_diffusion.sampling.sample_dpmpp_2m_sde", "k_diffusion.sampling.sample_dpmpp_2s_ancestral", "k_diffusion.sampling.sample_dpmpp_3m_sde", "k_diffusion.sampling.sample_dpmpp_sde", "k_diffusion.sampling.sample_euler", "k_diffusion.sampling.sample_euler_ancestral", "k_diffusion.sampling.sample_heun", "k_diffusion.sampling.sample_lms", "k_diffusion.sampling.to_d" ], "n_typable": 295, "n_typed": 0, "n_any": 0, "n_untyped": 295, "n_functions": 23, "n_function_overloads": 23, "n_function_params": 167, "n_methods": 17, "n_method_overloads": 17, "n_method_params": 73, "n_classes": 4, "n_attrs": 15, "n_properties": 0, "n_type_ignores": 0 }, { "path": "k_diffusion/utils.py", "symbol_reports": [ { "kind": "function", "name": "k_diffusion.utils.train_mode", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "k_diffusion.utils.FolderOfImages", "methods": [ { "kind": "function", "name": "FolderOfImages.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FolderOfImages.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FolderOfImages.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FolderOfImages.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "FolderOfImages.IMG_EXTENSIONS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FolderOfImages.root", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FolderOfImages.transform", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "FolderOfImages.paths", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 10, "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.utils.enable_stratified_accelerate", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.utils.rand_cosine_interpolated", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "k_diffusion.utils.get_safetensors_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "k_diffusion.utils.CSVLogger", "methods": [ { "kind": "function", "name": "CSVLogger.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CSVLogger.write", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CSVLogger.filename", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CSVLogger.columns", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CSVLogger.file", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.utils.tf32_mode", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.utils.n_params", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.utils.hf_datasets_augs_helper", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "k_diffusion.utils.ExponentialLR", "methods": [ { "kind": "function", "name": "ExponentialLR.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "ExponentialLR.get_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ExponentialLR._get_closed_form_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ExponentialLR.num_steps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ExponentialLR.decay", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ExponentialLR.warmup", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ExponentialLR.min_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 14, "n_typed": 0, "n_any": 0, "n_untyped": 14, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 7, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.utils.rand_log_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "k_diffusion.utils.EMAWarmup", "methods": [ { "kind": "function", "name": "EMAWarmup.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "EMAWarmup.state_dict", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EMAWarmup.load_state_dict", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "EMAWarmup.get_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EMAWarmup.step", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "EMAWarmup.inv_gamma", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EMAWarmup.power", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EMAWarmup.min_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EMAWarmup.max_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EMAWarmup.start_at", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EMAWarmup.last_epoch", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 18, "n_typed": 0, "n_any": 0, "n_untyped": 18, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 7, "n_attrs": 6, "n_properties": 0 }, { "kind": "class", "name": "k_diffusion.utils.InverseLR", "methods": [ { "kind": "function", "name": "InverseLR.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "InverseLR.get_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "InverseLR._get_closed_form_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "InverseLR.inv_gamma", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "InverseLR.power", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "InverseLR.warmup", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "InverseLR.min_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 14, "n_typed": 0, "n_any": 0, "n_untyped": 14, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 7, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.utils.rand_v_diffusion", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "k_diffusion.utils.ConstantLRWithWarmup", "methods": [ { "kind": "function", "name": "ConstantLRWithWarmup.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "ConstantLRWithWarmup.get_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ConstantLRWithWarmup._get_closed_form_lr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ConstantLRWithWarmup.warmup", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "k_diffusion.utils.download_file", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.utils.from_pil_image", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.utils.rand_log_logistic", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "k_diffusion.utils.append_dims", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "k_diffusion.utils.rand_split_log_normal", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "k_diffusion.utils.ema_update", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.utils.to_pil_image", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.utils.eval_mode", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "k_diffusion.utils.stratified_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.utils.stratified_with_settings", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "k_diffusion.utils.stratified_settings", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "k_diffusion.utils.rand_log_normal", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "k_diffusion.utils.enable_stratified", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "k_diffusion.utils.ema_update_dict", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "k_diffusion.utils", "names": [ "k_diffusion.utils.CSVLogger", "k_diffusion.utils.ConstantLRWithWarmup", "k_diffusion.utils.EMAWarmup", "k_diffusion.utils.ExponentialLR", "k_diffusion.utils.FolderOfImages", "k_diffusion.utils.InverseLR", "k_diffusion.utils.append_dims", "k_diffusion.utils.download_file", "k_diffusion.utils.ema_update", "k_diffusion.utils.ema_update_dict", "k_diffusion.utils.enable_stratified", "k_diffusion.utils.enable_stratified_accelerate", "k_diffusion.utils.eval_mode", "k_diffusion.utils.from_pil_image", "k_diffusion.utils.get_safetensors_metadata", "k_diffusion.utils.hf_datasets_augs_helper", "k_diffusion.utils.n_params", "k_diffusion.utils.rand_cosine_interpolated", "k_diffusion.utils.rand_log_logistic", "k_diffusion.utils.rand_log_normal", "k_diffusion.utils.rand_log_uniform", "k_diffusion.utils.rand_split_log_normal", "k_diffusion.utils.rand_v_diffusion", "k_diffusion.utils.stratified_settings", "k_diffusion.utils.stratified_uniform", "k_diffusion.utils.stratified_with_settings", "k_diffusion.utils.tf32_mode", "k_diffusion.utils.to_pil_image", "k_diffusion.utils.train_mode" ], "n_typable": 170, "n_typed": 0, "n_any": 0, "n_untyped": 170, "n_functions": 22, "n_function_overloads": 22, "n_function_params": 75, "n_methods": 20, "n_method_overloads": 20, "n_method_params": 31, "n_classes": 6, "n_attrs": 23, "n_properties": 0, "n_type_ignores": 0 } ], "typecheckers": {}, "n_modules": 12, "n_typable": 1156, "n_typed": 0, "n_any": 0, "n_untyped": 1156, "n_functions": 84, "n_function_overloads": 84, "n_function_params": 324, "n_methods": 167, "n_method_params": 410, "n_method_overloads": 167, "n_classes": 58, "n_attrs": 169, "n_properties": 3, "type_ignores": [], "n_type_ignores": 0 }