From 316114e660133b41756bcfe181aa43f048089a86 Mon Sep 17 00:00:00 2001 From: Bryce Date: Fri, 15 Dec 2023 12:31:28 -0800 Subject: [PATCH] docs: add docstrings Wrote an openai script and custom prompt to generate them. --- imaginairy/__init__.py | 1 + imaginairy/animations.py | 1 + imaginairy/api.py | 2 ++ imaginairy/api_refiners.py | 2 ++ imaginairy/cli/arg_schedule.py | 2 ++ imaginairy/cli/colorize.py | 2 ++ imaginairy/cli/describe.py | 2 ++ imaginairy/cli/edit.py | 2 ++ imaginairy/cli/edit_demo.py | 2 ++ imaginairy/cli/imagine.py | 2 ++ imaginairy/cli/main.py | 2 ++ imaginairy/cli/run_api.py | 2 ++ imaginairy/cli/shared.py | 2 ++ imaginairy/cli/train.py | 2 ++ imaginairy/cli/upscale.py | 2 ++ imaginairy/cli/videogen.py | 2 ++ imaginairy/colorize.py | 2 ++ imaginairy/config.py | 2 ++ imaginairy/enhancers/blur_detect.py | 2 ++ imaginairy/enhancers/clip_masking.py | 2 ++ imaginairy/enhancers/describe_image_blip.py | 2 ++ imaginairy/enhancers/describe_image_clip.py | 2 ++ imaginairy/enhancers/face_restoration_codeformer.py | 2 ++ imaginairy/enhancers/facecrop.py | 2 ++ imaginairy/enhancers/prompt_expansion.py | 2 ++ imaginairy/enhancers/upscale_realesrgan.py | 2 ++ imaginairy/enhancers/upscale_riverwing.py | 2 ++ imaginairy/feather_tile.py | 2 ++ imaginairy/http_app/app.py | 2 ++ imaginairy/http_app/stablestudio/models.py | 2 ++ imaginairy/http_app/stablestudio/routes.py | 2 ++ imaginairy/http_app/utils.py | 2 ++ imaginairy/img_processors/hed_boundary.py | 2 ++ imaginairy/img_processors/openpose.py | 2 ++ imaginairy/log_utils.py | 2 ++ imaginairy/lr_scheduler.py | 2 ++ imaginairy/model_manager.py | 2 ++ imaginairy/modules/attention.py | 2 ++ imaginairy/modules/autoencoder.py | 2 ++ imaginairy/modules/cldm.py | 2 ++ imaginairy/modules/clip_embedders.py | 2 ++ imaginairy/modules/diffusion/model.py | 2 ++ imaginairy/modules/diffusion/openaimodel.py | 2 ++ imaginairy/modules/diffusion/upscaling.py | 2 ++ imaginairy/modules/diffusion/util.py | 2 ++ imaginairy/modules/distributions.py | 2 ++ imaginairy/modules/ema.py | 2 ++ imaginairy/modules/encoders.py | 2 ++ imaginairy/modules/midas/api.py | 2 ++ imaginairy/modules/midas/midas/base_model.py | 2 ++ imaginairy/modules/midas/midas/blocks.py | 2 ++ imaginairy/modules/midas/midas/dpt_depth.py | 2 ++ imaginairy/modules/midas/midas/transforms.py | 2 ++ imaginairy/modules/midas/midas/vit.py | 2 ++ imaginairy/modules/refiners_sd.py | 2 ++ imaginairy/modules/sgm/attention.py | 2 ++ imaginairy/modules/sgm/autoencoder.py | 2 ++ .../modules/sgm/autoencoding/losses/discriminator_loss.py | 2 ++ imaginairy/modules/sgm/autoencoding/losses/lpips.py | 2 ++ imaginairy/modules/sgm/autoencoding/lpips/model/model.py | 2 ++ imaginairy/modules/sgm/autoencoding/lpips/util.py | 2 ++ imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py | 2 ++ imaginairy/modules/sgm/autoencoding/regularizers/base.py | 2 ++ imaginairy/modules/sgm/autoencoding/regularizers/quantize.py | 2 ++ imaginairy/modules/sgm/autoencoding/temporal_ae.py | 2 ++ imaginairy/modules/sgm/diffusion.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/denoiser.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/denoiser_scaling.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/discretizer.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/guiders.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/loss.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/loss_weighting.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/model.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/openaimodel.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/sampling_utils.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/sigma_sampling.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/video_model.py | 2 ++ imaginairy/modules/sgm/diffusionmodules/wrappers.py | 2 ++ imaginairy/modules/sgm/distributions/distributions.py | 2 ++ imaginairy/modules/sgm/encoders/modules.py | 2 ++ imaginairy/modules/sgm/video_attention.py | 2 ++ imaginairy/outpaint.py | 2 ++ imaginairy/paths.py | 2 ++ imaginairy/prompt_schedules.py | 2 ++ imaginairy/roi_utils.py | 2 ++ imaginairy/safety.py | 2 ++ imaginairy/samplers/base.py | 2 ++ imaginairy/samplers/ddim.py | 2 ++ imaginairy/samplers/kdiff.py | 2 ++ imaginairy/samplers/plms.py | 2 ++ imaginairy/schema.py | 2 ++ imaginairy/surprise_me.py | 2 ++ imaginairy/train.py | 2 ++ imaginairy/training_tools/image_prep.py | 2 ++ imaginairy/training_tools/prune_model.py | 2 ++ imaginairy/training_tools/single_concept.py | 2 ++ imaginairy/utils/data_distorter.py | 2 ++ imaginairy/utils/debug_info.py | 2 ++ imaginairy/utils/model_cache.py | 2 ++ imaginairy/utils/named_resolutions.py | 2 ++ imaginairy/utils/torch_installer.py | 2 ++ imaginairy/vendored/basicsr/arch_util.py | 2 ++ imaginairy/vendored/basicsr/img_util.py | 2 ++ imaginairy/vendored/basicsr/rrdbnet_arch.py | 2 ++ imaginairy/vendored/blip/blip_itm.py | 2 ++ imaginairy/vendored/blip/blip_nlvr.py | 2 ++ imaginairy/vendored/blip/blip_retrieval.py | 2 ++ imaginairy/vendored/blip/blip_vqa.py | 2 ++ imaginairy/vendored/blip/nlvr_encoder.py | 2 ++ imaginairy/vendored/clip/clip.py | 2 ++ imaginairy/vendored/clip/model.py | 2 ++ imaginairy/vendored/clip/simple_tokenizer.py | 2 ++ imaginairy/vendored/clip/version.py | 1 + imaginairy/vendored/codeformer/codeformer_arch.py | 2 ++ imaginairy/version.py | 3 +++ imaginairy/video_sample.py | 2 ++ imaginairy/weight_management/conversion.py | 2 ++ imaginairy/weight_management/execution_trace.py | 2 ++ imaginairy/weight_management/generate_conversion_maps.py | 2 ++ imaginairy/weight_management/generate_weight_info.py | 2 ++ imaginairy/weight_management/pattern_collapse.py | 3 +++ imaginairy/weight_management/utils.py | 2 ++ 123 files changed, 245 insertions(+) diff --git a/imaginairy/__init__.py b/imaginairy/__init__.py index aede5c9..69425f1 100644 --- a/imaginairy/__init__.py +++ b/imaginairy/__init__.py @@ -1,3 +1,4 @@ +"""Pythonic AI generation of images and videos""" import os # tells pytorch to allow MPS usage (for Mac M1 compatibility) diff --git a/imaginairy/animations.py b/imaginairy/animations.py index 9afe2a6..17c4b79 100644 --- a/imaginairy/animations.py +++ b/imaginairy/animations.py @@ -1,3 +1,4 @@ +"""Functions for creating animations from images.""" import os.path import cv2 diff --git a/imaginairy/api.py b/imaginairy/api.py index d98f434..e7c5f80 100755 --- a/imaginairy/api.py +++ b/imaginairy/api.py @@ -1,3 +1,5 @@ +"""Functions for generating and processing images""" + import logging import os import re diff --git a/imaginairy/api_refiners.py b/imaginairy/api_refiners.py index 4599611..38aba26 100644 --- a/imaginairy/api_refiners.py +++ b/imaginairy/api_refiners.py @@ -1,3 +1,5 @@ +"""Functions for generating refined images""" + import logging from typing import List, Optional diff --git a/imaginairy/cli/arg_schedule.py b/imaginairy/cli/arg_schedule.py index 1569462..a9e5a88 100644 --- a/imaginairy/cli/arg_schedule.py +++ b/imaginairy/cli/arg_schedule.py @@ -1,3 +1,5 @@ +"""Decorator and parsers for command scheduling""" + from typing import Iterable from imaginairy.utils import frange diff --git a/imaginairy/cli/colorize.py b/imaginairy/cli/colorize.py index 156b0cf..23d6ba9 100644 --- a/imaginairy/cli/colorize.py +++ b/imaginairy/cli/colorize.py @@ -1,3 +1,5 @@ +"""Command for colorizing images with AI""" + import logging import click diff --git a/imaginairy/cli/describe.py b/imaginairy/cli/describe.py index 78ab9d1..30d36c7 100644 --- a/imaginairy/cli/describe.py +++ b/imaginairy/cli/describe.py @@ -1,3 +1,5 @@ +"""Command for generating image descriptions""" + import click diff --git a/imaginairy/cli/edit.py b/imaginairy/cli/edit.py index c25120b..91be7fd 100644 --- a/imaginairy/cli/edit.py +++ b/imaginairy/cli/edit.py @@ -1,3 +1,5 @@ +"""Command for AI-powered image editing""" + import click from imaginairy import config diff --git a/imaginairy/cli/edit_demo.py b/imaginairy/cli/edit_demo.py index 94282f6..aa32fd1 100644 --- a/imaginairy/cli/edit_demo.py +++ b/imaginairy/cli/edit_demo.py @@ -1,3 +1,5 @@ +"""Command for applying preset image edits""" + import click diff --git a/imaginairy/cli/imagine.py b/imaginairy/cli/imagine.py index 7ee3b62..0e44e7d 100644 --- a/imaginairy/cli/imagine.py +++ b/imaginairy/cli/imagine.py @@ -1,3 +1,5 @@ +"""Command-line interface for AI-driven image generation""" + import click from imaginairy.cli.clickshell_mod import ImagineColorsCommand diff --git a/imaginairy/cli/main.py b/imaginairy/cli/main.py index aab7358..9996dd5 100644 --- a/imaginairy/cli/main.py +++ b/imaginairy/cli/main.py @@ -1,3 +1,5 @@ +"""CLI for AI-powered image generation""" + import logging import click diff --git a/imaginairy/cli/run_api.py b/imaginairy/cli/run_api.py index 10d5e77..b01b52d 100644 --- a/imaginairy/cli/run_api.py +++ b/imaginairy/cli/run_api.py @@ -1,3 +1,5 @@ +"""Code for starting an HTTP API server""" + import logging import click diff --git a/imaginairy/cli/shared.py b/imaginairy/cli/shared.py index a7f0283..ff7f2e3 100644 --- a/imaginairy/cli/shared.py +++ b/imaginairy/cli/shared.py @@ -1,3 +1,5 @@ +"""Context managers and functions for image generation CLI""" + import logging import math from contextlib import contextmanager diff --git a/imaginairy/cli/train.py b/imaginairy/cli/train.py index b952f14..d23747f 100644 --- a/imaginairy/cli/train.py +++ b/imaginairy/cli/train.py @@ -1,3 +1,5 @@ +"""CLI commands for model training and image preparation""" + import logging import click diff --git a/imaginairy/cli/upscale.py b/imaginairy/cli/upscale.py index 6494513..5894e36 100644 --- a/imaginairy/cli/upscale.py +++ b/imaginairy/cli/upscale.py @@ -1,3 +1,5 @@ +"""Command for upscaling images with AI""" + import logging import click diff --git a/imaginairy/cli/videogen.py b/imaginairy/cli/videogen.py index 1a5700b..d8e19ea 100644 --- a/imaginairy/cli/videogen.py +++ b/imaginairy/cli/videogen.py @@ -1,3 +1,5 @@ +"""Command for generating AI-powered videos""" + import logging import click diff --git a/imaginairy/colorize.py b/imaginairy/colorize.py index f50f805..554d6cf 100644 --- a/imaginairy/colorize.py +++ b/imaginairy/colorize.py @@ -1,3 +1,5 @@ +"""Functions for colorizing black and white images""" + import logging from PIL import Image, ImageEnhance, ImageStat diff --git a/imaginairy/config.py b/imaginairy/config.py index b8934fd..115821f 100644 --- a/imaginairy/config.py +++ b/imaginairy/config.py @@ -1,3 +1,5 @@ +"""Classes and constants for AI model configuration""" + from dataclasses import dataclass from typing import Any, List diff --git a/imaginairy/enhancers/blur_detect.py b/imaginairy/enhancers/blur_detect.py index 2a0beff..1647db8 100644 --- a/imaginairy/enhancers/blur_detect.py +++ b/imaginairy/enhancers/blur_detect.py @@ -1,3 +1,5 @@ +"""Functions for assessing image blurriness""" + import cv2 from imaginairy.img_utils import pillow_img_to_opencv_img diff --git a/imaginairy/enhancers/clip_masking.py b/imaginairy/enhancers/clip_masking.py index 9eb31bd..7ea94ac 100644 --- a/imaginairy/enhancers/clip_masking.py +++ b/imaginairy/enhancers/clip_masking.py @@ -1,3 +1,5 @@ +"""Functions for generating and processing image masks""" + from functools import lru_cache from typing import Optional, Sequence diff --git a/imaginairy/enhancers/describe_image_blip.py b/imaginairy/enhancers/describe_image_blip.py index 12fa6e9..55f948b 100644 --- a/imaginairy/enhancers/describe_image_blip.py +++ b/imaginairy/enhancers/describe_image_blip.py @@ -1,3 +1,5 @@ +"""Functions for generating image captions""" + import os import os.path from functools import lru_cache diff --git a/imaginairy/enhancers/describe_image_clip.py b/imaginairy/enhancers/describe_image_clip.py index 42cec54..c40482e 100644 --- a/imaginairy/enhancers/describe_image_clip.py +++ b/imaginairy/enhancers/describe_image_clip.py @@ -1,3 +1,5 @@ +"""Functions for image-text similarity assessment""" + from functools import lru_cache from typing import Sequence diff --git a/imaginairy/enhancers/face_restoration_codeformer.py b/imaginairy/enhancers/face_restoration_codeformer.py index 377139c..1f48847 100644 --- a/imaginairy/enhancers/face_restoration_codeformer.py +++ b/imaginairy/enhancers/face_restoration_codeformer.py @@ -1,3 +1,5 @@ +"""Code for enhancing facial images""" + import logging from functools import lru_cache diff --git a/imaginairy/enhancers/facecrop.py b/imaginairy/enhancers/facecrop.py index e22ece4..cbf295b 100644 --- a/imaginairy/enhancers/facecrop.py +++ b/imaginairy/enhancers/facecrop.py @@ -1,3 +1,5 @@ +"""Functions for detecting and cropping faces""" + import numpy as np from imaginairy.enhancers.face_restoration_codeformer import face_restore_helper diff --git a/imaginairy/enhancers/prompt_expansion.py b/imaginairy/enhancers/prompt_expansion.py index 044c636..0741e47 100644 --- a/imaginairy/enhancers/prompt_expansion.py +++ b/imaginairy/enhancers/prompt_expansion.py @@ -1,3 +1,5 @@ +"""Functions for expanding text prompts with phraselists""" + import gzip import os.path import random diff --git a/imaginairy/enhancers/upscale_realesrgan.py b/imaginairy/enhancers/upscale_realesrgan.py index c564216..5050788 100644 --- a/imaginairy/enhancers/upscale_realesrgan.py +++ b/imaginairy/enhancers/upscale_realesrgan.py @@ -1,3 +1,5 @@ +"""Functions for image upscaling using RealESRGAN""" + import numpy as np import torch from PIL import Image diff --git a/imaginairy/enhancers/upscale_riverwing.py b/imaginairy/enhancers/upscale_riverwing.py index 81c43a8..d0a975f 100644 --- a/imaginairy/enhancers/upscale_riverwing.py +++ b/imaginairy/enhancers/upscale_riverwing.py @@ -1,3 +1,5 @@ +"""Classes and functions for image upscaling""" + from functools import lru_cache import numpy as np diff --git a/imaginairy/feather_tile.py b/imaginairy/feather_tile.py index 7c154fb..c06b286 100644 --- a/imaginairy/feather_tile.py +++ b/imaginairy/feather_tile.py @@ -1,3 +1,5 @@ +"""Functions for image tiling and reconstruction""" + # inspired by https://github.com/ProGamerGov/neural-dream/blob/master/neural_dream/dream_tile.py # but with all the bugs fixed and lots of simplifications # MIT License diff --git a/imaginairy/http_app/app.py b/imaginairy/http_app/app.py index 490b8ce..46dfb44 100644 --- a/imaginairy/http_app/app.py +++ b/imaginairy/http_app/app.py @@ -1,3 +1,5 @@ +"""FastAPI application for image generation""" + import logging import os.path import sys diff --git a/imaginairy/http_app/stablestudio/models.py b/imaginairy/http_app/stablestudio/models.py index a0fb297..0cbd293 100644 --- a/imaginairy/http_app/stablestudio/models.py +++ b/imaginairy/http_app/stablestudio/models.py @@ -1,3 +1,5 @@ +"""Classes for image generation API models""" + from datetime import datetime from typing import List, Optional diff --git a/imaginairy/http_app/stablestudio/routes.py b/imaginairy/http_app/stablestudio/routes.py index bebe5e2..a5c3192 100644 --- a/imaginairy/http_app/stablestudio/routes.py +++ b/imaginairy/http_app/stablestudio/routes.py @@ -1,3 +1,5 @@ +"""API routes for image generation service""" + import uuid from fastapi import APIRouter diff --git a/imaginairy/http_app/utils.py b/imaginairy/http_app/utils.py index 2eb7618..a416fe5 100644 --- a/imaginairy/http_app/utils.py +++ b/imaginairy/http_app/utils.py @@ -1,3 +1,5 @@ +"""Functions and class for image generation and encoding""" + import base64 from io import BytesIO diff --git a/imaginairy/img_processors/hed_boundary.py b/imaginairy/img_processors/hed_boundary.py index 3099b5c..6da7c0e 100644 --- a/imaginairy/img_processors/hed_boundary.py +++ b/imaginairy/img_processors/hed_boundary.py @@ -1,3 +1,5 @@ +"""Classes and functions for edge detection in images""" + from functools import lru_cache import cv2 diff --git a/imaginairy/img_processors/openpose.py b/imaginairy/img_processors/openpose.py index 9ff4ed7..81e1444 100644 --- a/imaginairy/img_processors/openpose.py +++ b/imaginairy/img_processors/openpose.py @@ -1,3 +1,5 @@ +"""Functions for human pose estimation""" + import math from collections import OrderedDict from functools import lru_cache diff --git a/imaginairy/log_utils.py b/imaginairy/log_utils.py index 0ef7d36..1a46f90 100644 --- a/imaginairy/log_utils.py +++ b/imaginairy/log_utils.py @@ -1,3 +1,5 @@ +"""Utilities for image generation logging""" + import logging import logging.config import re diff --git a/imaginairy/lr_scheduler.py b/imaginairy/lr_scheduler.py index f36ff9e..3deb0ec 100644 --- a/imaginairy/lr_scheduler.py +++ b/imaginairy/lr_scheduler.py @@ -1,3 +1,5 @@ +"""Classes for learning rate scheduling""" + import numpy as np diff --git a/imaginairy/model_manager.py b/imaginairy/model_manager.py index 80dcee6..505c0d6 100644 --- a/imaginairy/model_manager.py +++ b/imaginairy/model_manager.py @@ -1,3 +1,5 @@ +"""Classes and functions for managing AI models""" + import logging import os import re diff --git a/imaginairy/modules/attention.py b/imaginairy/modules/attention.py index cf409a2..f5b9e7f 100644 --- a/imaginairy/modules/attention.py +++ b/imaginairy/modules/attention.py @@ -1,3 +1,5 @@ +"""Classes and functions for attention mechanisms in neural networks""" + import math from functools import lru_cache diff --git a/imaginairy/modules/autoencoder.py b/imaginairy/modules/autoencoder.py index b932a02..ac42202 100644 --- a/imaginairy/modules/autoencoder.py +++ b/imaginairy/modules/autoencoder.py @@ -1,3 +1,5 @@ +"""Classes for image autoencoding and manipulation""" + # pylama:ignore=W0613 import logging import math diff --git a/imaginairy/modules/cldm.py b/imaginairy/modules/cldm.py index fc77df3..ed32613 100644 --- a/imaginairy/modules/cldm.py +++ b/imaginairy/modules/cldm.py @@ -1,3 +1,5 @@ +"""Classes for controlled image generation""" + import torch from torch import nn diff --git a/imaginairy/modules/clip_embedders.py b/imaginairy/modules/clip_embedders.py index 4f7da44..50e2446 100644 --- a/imaginairy/modules/clip_embedders.py +++ b/imaginairy/modules/clip_embedders.py @@ -1,3 +1,5 @@ +"""Classes for text and image encoding""" + import kornia import torch from einops import repeat diff --git a/imaginairy/modules/diffusion/model.py b/imaginairy/modules/diffusion/model.py index 89053e6..4dc0234 100644 --- a/imaginairy/modules/diffusion/model.py +++ b/imaginairy/modules/diffusion/model.py @@ -1,3 +1,5 @@ +"""Classes for diffusion-based image generation""" + # pytorch_diffusion + derived encoder decoder import gc import math diff --git a/imaginairy/modules/diffusion/openaimodel.py b/imaginairy/modules/diffusion/openaimodel.py index 12e3b09..42bf33c 100644 --- a/imaginairy/modules/diffusion/openaimodel.py +++ b/imaginairy/modules/diffusion/openaimodel.py @@ -1,3 +1,5 @@ +"""Classes for generative diffusion models""" + import math from abc import abstractmethod from typing import Optional diff --git a/imaginairy/modules/diffusion/upscaling.py b/imaginairy/modules/diffusion/upscaling.py index 06103c0..c889216 100644 --- a/imaginairy/modules/diffusion/upscaling.py +++ b/imaginairy/modules/diffusion/upscaling.py @@ -1,3 +1,5 @@ +"""Classes for image upscaling diffusion models""" + from functools import partial import numpy as np diff --git a/imaginairy/modules/diffusion/util.py b/imaginairy/modules/diffusion/util.py index 6893e46..2fc02f5 100644 --- a/imaginairy/modules/diffusion/util.py +++ b/imaginairy/modules/diffusion/util.py @@ -1,3 +1,5 @@ +"""Utilities for diffusion model operations""" + # adopted from # https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py # and diff --git a/imaginairy/modules/distributions.py b/imaginairy/modules/distributions.py index 62a671b..f2d54e2 100644 --- a/imaginairy/modules/distributions.py +++ b/imaginairy/modules/distributions.py @@ -1,3 +1,5 @@ +"""Classes for parametric Gaussian distributions""" + import numpy as np import torch diff --git a/imaginairy/modules/ema.py b/imaginairy/modules/ema.py index 68a164e..a80c539 100644 --- a/imaginairy/modules/ema.py +++ b/imaginairy/modules/ema.py @@ -1,3 +1,5 @@ +"""Classes for exponential moving average tracking""" + import torch from torch import nn diff --git a/imaginairy/modules/encoders.py b/imaginairy/modules/encoders.py index a16cb0b..aaaa88e 100644 --- a/imaginairy/modules/encoders.py +++ b/imaginairy/modules/encoders.py @@ -1,3 +1,5 @@ +"""Classes for text encoding transformations""" + import open_clip import torch from torch import nn diff --git a/imaginairy/modules/midas/api.py b/imaginairy/modules/midas/api.py index 216bfbb..f2cda81 100644 --- a/imaginairy/modules/midas/api.py +++ b/imaginairy/modules/midas/api.py @@ -1,3 +1,5 @@ +"""Code for depth estimation with MiDaS models""" + # based on https://github.com/isl-org/MiDaS from functools import lru_cache diff --git a/imaginairy/modules/midas/midas/base_model.py b/imaginairy/modules/midas/midas/base_model.py index dad86d1..d90399f 100644 --- a/imaginairy/modules/midas/midas/base_model.py +++ b/imaginairy/modules/midas/midas/base_model.py @@ -1,3 +1,5 @@ +"""Classes for loading model weights""" + import torch from imaginairy import config diff --git a/imaginairy/modules/midas/midas/blocks.py b/imaginairy/modules/midas/midas/blocks.py index 1348457..065cba1 100644 --- a/imaginairy/modules/midas/midas/blocks.py +++ b/imaginairy/modules/midas/midas/blocks.py @@ -1,3 +1,5 @@ +"""Functions and classes for neural network construction""" + import torch from torch import nn diff --git a/imaginairy/modules/midas/midas/dpt_depth.py b/imaginairy/modules/midas/midas/dpt_depth.py index f7d45b7..8de11bd 100644 --- a/imaginairy/modules/midas/midas/dpt_depth.py +++ b/imaginairy/modules/midas/midas/dpt_depth.py @@ -1,3 +1,5 @@ +"""Classes for depth estimation from images""" + import torch from torch import nn diff --git a/imaginairy/modules/midas/midas/transforms.py b/imaginairy/modules/midas/midas/transforms.py index 1b7b13f..9659f23 100644 --- a/imaginairy/modules/midas/midas/transforms.py +++ b/imaginairy/modules/midas/midas/transforms.py @@ -1,3 +1,5 @@ +"""Functions and classes for image preprocessing""" + import math import cv2 diff --git a/imaginairy/modules/midas/midas/vit.py b/imaginairy/modules/midas/midas/vit.py index b799c28..841164b 100644 --- a/imaginairy/modules/midas/midas/vit.py +++ b/imaginairy/modules/midas/midas/vit.py @@ -1,3 +1,5 @@ +"""Classes and functions for Vision Transformer processing""" + import math import types diff --git a/imaginairy/modules/refiners_sd.py b/imaginairy/modules/refiners_sd.py index 86a37c9..b7ed317 100644 --- a/imaginairy/modules/refiners_sd.py +++ b/imaginairy/modules/refiners_sd.py @@ -1,3 +1,5 @@ +"""Refinement modules for image generation""" + import logging import math from functools import lru_cache diff --git a/imaginairy/modules/sgm/attention.py b/imaginairy/modules/sgm/attention.py index 00acd08..069fe7d 100644 --- a/imaginairy/modules/sgm/attention.py +++ b/imaginairy/modules/sgm/attention.py @@ -1,3 +1,5 @@ +"""Classes for attention-based neural networks""" + import logging import math from inspect import isfunction diff --git a/imaginairy/modules/sgm/autoencoder.py b/imaginairy/modules/sgm/autoencoder.py index f7bc86f..4f5cac1 100644 --- a/imaginairy/modules/sgm/autoencoder.py +++ b/imaginairy/modules/sgm/autoencoder.py @@ -1,3 +1,5 @@ +"""Classes for autoencoder-based image generation""" + import logging import math import re diff --git a/imaginairy/modules/sgm/autoencoding/losses/discriminator_loss.py b/imaginairy/modules/sgm/autoencoding/losses/discriminator_loss.py index 499d90f..b14f5f9 100644 --- a/imaginairy/modules/sgm/autoencoding/losses/discriminator_loss.py +++ b/imaginairy/modules/sgm/autoencoding/losses/discriminator_loss.py @@ -1,3 +1,5 @@ +"""Classes for discriminator loss computation""" + from typing import Dict, Iterator, List, Optional, Tuple, Union import numpy as np diff --git a/imaginairy/modules/sgm/autoencoding/losses/lpips.py b/imaginairy/modules/sgm/autoencoding/losses/lpips.py index c3a7004..cd9433a 100644 --- a/imaginairy/modules/sgm/autoencoding/losses/lpips.py +++ b/imaginairy/modules/sgm/autoencoding/losses/lpips.py @@ -1,3 +1,5 @@ +"""Classes for latent space perceptual loss""" + import torch import torch.nn as nn diff --git a/imaginairy/modules/sgm/autoencoding/lpips/model/model.py b/imaginairy/modules/sgm/autoencoding/lpips/model/model.py index 1f69b5c..52607fd 100644 --- a/imaginairy/modules/sgm/autoencoding/lpips/model/model.py +++ b/imaginairy/modules/sgm/autoencoding/lpips/model/model.py @@ -1,3 +1,5 @@ +"""Classes for image discrimination""" + import functools import torch.nn as nn diff --git a/imaginairy/modules/sgm/autoencoding/lpips/util.py b/imaginairy/modules/sgm/autoencoding/lpips/util.py index 06d1bd8..13e2a53 100644 --- a/imaginairy/modules/sgm/autoencoding/lpips/util.py +++ b/imaginairy/modules/sgm/autoencoding/lpips/util.py @@ -1,3 +1,5 @@ +"""Utilities for model checkpoint management and normalization layers""" + import hashlib import os diff --git a/imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py b/imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py index 6195f0a..8e0720a 100644 --- a/imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py +++ b/imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py @@ -1,3 +1,5 @@ +"""Functions for discriminator loss calculation""" + import torch import torch.nn.functional as F diff --git a/imaginairy/modules/sgm/autoencoding/regularizers/base.py b/imaginairy/modules/sgm/autoencoding/regularizers/base.py index e833a10..f572d51 100644 --- a/imaginairy/modules/sgm/autoencoding/regularizers/base.py +++ b/imaginairy/modules/sgm/autoencoding/regularizers/base.py @@ -1,3 +1,5 @@ +"""Classes and functions for tensor regularization""" + from abc import abstractmethod from typing import Any, Tuple diff --git a/imaginairy/modules/sgm/autoencoding/regularizers/quantize.py b/imaginairy/modules/sgm/autoencoding/regularizers/quantize.py index 3113076..cebfcca 100644 --- a/imaginairy/modules/sgm/autoencoding/regularizers/quantize.py +++ b/imaginairy/modules/sgm/autoencoding/regularizers/quantize.py @@ -1,3 +1,5 @@ +"""Classes for vector quantization regularization""" + import logging from abc import abstractmethod from typing import Dict, Iterator, Literal, Optional, Tuple, Union diff --git a/imaginairy/modules/sgm/autoencoding/temporal_ae.py b/imaginairy/modules/sgm/autoencoding/temporal_ae.py index 6207d78..b03afbb 100644 --- a/imaginairy/modules/sgm/autoencoding/temporal_ae.py +++ b/imaginairy/modules/sgm/autoencoding/temporal_ae.py @@ -1,3 +1,5 @@ +"""Classes for spatio-temporal video processing""" + import logging from typing import Callable, Iterable, Union diff --git a/imaginairy/modules/sgm/diffusion.py b/imaginairy/modules/sgm/diffusion.py index ee49bc3..71301df 100644 --- a/imaginairy/modules/sgm/diffusion.py +++ b/imaginairy/modules/sgm/diffusion.py @@ -1,3 +1,5 @@ +"""Classes for stochastic diffusion models""" + import logging import math from contextlib import contextmanager diff --git a/imaginairy/modules/sgm/diffusionmodules/denoiser.py b/imaginairy/modules/sgm/diffusionmodules/denoiser.py index e4e9b5c..1196684 100644 --- a/imaginairy/modules/sgm/diffusionmodules/denoiser.py +++ b/imaginairy/modules/sgm/diffusionmodules/denoiser.py @@ -1,3 +1,5 @@ +"""Classes for image denoising operations""" + from typing import TYPE_CHECKING, Dict, Union import torch diff --git a/imaginairy/modules/sgm/diffusionmodules/denoiser_scaling.py b/imaginairy/modules/sgm/diffusionmodules/denoiser_scaling.py index f4e287b..4e9fba1 100644 --- a/imaginairy/modules/sgm/diffusionmodules/denoiser_scaling.py +++ b/imaginairy/modules/sgm/diffusionmodules/denoiser_scaling.py @@ -1,3 +1,5 @@ +"""Classes for denoiser scaling calculations""" + from abc import ABC, abstractmethod from typing import Tuple diff --git a/imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py b/imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py index b8b03ca..49ff10d 100644 --- a/imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py +++ b/imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py @@ -1,3 +1,5 @@ +"""Classes for denoiser weighting strategies""" + import torch diff --git a/imaginairy/modules/sgm/diffusionmodules/discretizer.py b/imaginairy/modules/sgm/diffusionmodules/discretizer.py index d6fc6dc..e42916a 100644 --- a/imaginairy/modules/sgm/diffusionmodules/discretizer.py +++ b/imaginairy/modules/sgm/diffusionmodules/discretizer.py @@ -1,3 +1,5 @@ +"""Classes for discretizing diffusion steps""" + from abc import abstractmethod from functools import partial diff --git a/imaginairy/modules/sgm/diffusionmodules/guiders.py b/imaginairy/modules/sgm/diffusionmodules/guiders.py index 66a38a2..8d813a3 100644 --- a/imaginairy/modules/sgm/diffusionmodules/guiders.py +++ b/imaginairy/modules/sgm/diffusionmodules/guiders.py @@ -1,3 +1,5 @@ +"""Classes for guiding diffusion models""" + import logging from abc import ABC, abstractmethod from typing import Dict, List, Optional, Tuple, Union diff --git a/imaginairy/modules/sgm/diffusionmodules/loss.py b/imaginairy/modules/sgm/diffusionmodules/loss.py index 0ea9e4f..8e100b3 100644 --- a/imaginairy/modules/sgm/diffusionmodules/loss.py +++ b/imaginairy/modules/sgm/diffusionmodules/loss.py @@ -1,3 +1,5 @@ +"""Classes for diffusion model loss calculation""" + from typing import Dict, List, Optional, Tuple, Union import torch diff --git a/imaginairy/modules/sgm/diffusionmodules/loss_weighting.py b/imaginairy/modules/sgm/diffusionmodules/loss_weighting.py index e12c0a7..6cf8eab 100644 --- a/imaginairy/modules/sgm/diffusionmodules/loss_weighting.py +++ b/imaginairy/modules/sgm/diffusionmodules/loss_weighting.py @@ -1,3 +1,5 @@ +"""Classes for diffusion loss weighting""" + from abc import ABC, abstractmethod import torch diff --git a/imaginairy/modules/sgm/diffusionmodules/model.py b/imaginairy/modules/sgm/diffusionmodules/model.py index 0fc0dea..05419ad 100644 --- a/imaginairy/modules/sgm/diffusionmodules/model.py +++ b/imaginairy/modules/sgm/diffusionmodules/model.py @@ -1,3 +1,5 @@ +"""Classes for diffusion-based image generation""" + # pytorch_diffusion + derived encoder decoder import logging import math diff --git a/imaginairy/modules/sgm/diffusionmodules/openaimodel.py b/imaginairy/modules/sgm/diffusionmodules/openaimodel.py index 15c414d..d35423a 100644 --- a/imaginairy/modules/sgm/diffusionmodules/openaimodel.py +++ b/imaginairy/modules/sgm/diffusionmodules/openaimodel.py @@ -1,3 +1,5 @@ +"""Classes for generative image modeling""" + import logging import math from abc import abstractmethod diff --git a/imaginairy/modules/sgm/diffusionmodules/sampling_utils.py b/imaginairy/modules/sgm/diffusionmodules/sampling_utils.py index 280c21d..33ef04e 100644 --- a/imaginairy/modules/sgm/diffusionmodules/sampling_utils.py +++ b/imaginairy/modules/sgm/diffusionmodules/sampling_utils.py @@ -1,3 +1,5 @@ +"""Functions for diffusion sampling calculations""" + import torch from scipy import integrate diff --git a/imaginairy/modules/sgm/diffusionmodules/sigma_sampling.py b/imaginairy/modules/sgm/diffusionmodules/sigma_sampling.py index 30d93b0..68524b3 100644 --- a/imaginairy/modules/sgm/diffusionmodules/sigma_sampling.py +++ b/imaginairy/modules/sgm/diffusionmodules/sigma_sampling.py @@ -1,3 +1,5 @@ +"""Classes for stochastic sampling strategies""" + import torch from imaginairy.utils import default, instantiate_from_config diff --git a/imaginairy/modules/sgm/diffusionmodules/video_model.py b/imaginairy/modules/sgm/diffusionmodules/video_model.py index 55b1aab..249eab9 100644 --- a/imaginairy/modules/sgm/diffusionmodules/video_model.py +++ b/imaginairy/modules/sgm/diffusionmodules/video_model.py @@ -1,3 +1,5 @@ +"""Classes for video generation models""" + from typing import List, Optional, Union import torch as th diff --git a/imaginairy/modules/sgm/diffusionmodules/wrappers.py b/imaginairy/modules/sgm/diffusionmodules/wrappers.py index 849a1a2..61e04b0 100644 --- a/imaginairy/modules/sgm/diffusionmodules/wrappers.py +++ b/imaginairy/modules/sgm/diffusionmodules/wrappers.py @@ -1,3 +1,5 @@ +"""Wrappers for diffusion model integration""" + import torch import torch.nn as nn from packaging import version diff --git a/imaginairy/modules/sgm/distributions/distributions.py b/imaginairy/modules/sgm/distributions/distributions.py index 4300035..9527c48 100644 --- a/imaginairy/modules/sgm/distributions/distributions.py +++ b/imaginairy/modules/sgm/distributions/distributions.py @@ -1,3 +1,5 @@ +"""Classes for probabilistic distributions handling""" + import numpy as np import torch diff --git a/imaginairy/modules/sgm/encoders/modules.py b/imaginairy/modules/sgm/encoders/modules.py index 0228537..a308641 100644 --- a/imaginairy/modules/sgm/encoders/modules.py +++ b/imaginairy/modules/sgm/encoders/modules.py @@ -1,3 +1,5 @@ +"""Classes for image and text encoding""" + import math from contextlib import nullcontext from functools import partial diff --git a/imaginairy/modules/sgm/video_attention.py b/imaginairy/modules/sgm/video_attention.py index 3221649..765eccd 100644 --- a/imaginairy/modules/sgm/video_attention.py +++ b/imaginairy/modules/sgm/video_attention.py @@ -1,3 +1,5 @@ +"""Classes for video sequence transformation""" + import logging from typing import Optional diff --git a/imaginairy/outpaint.py b/imaginairy/outpaint.py index 8f6fab0..1b4ba3d 100644 --- a/imaginairy/outpaint.py +++ b/imaginairy/outpaint.py @@ -1,3 +1,5 @@ +"""Functions for image outpainting preparation""" + import re import torch diff --git a/imaginairy/paths.py b/imaginairy/paths.py index 5a80639..6e13eae 100644 --- a/imaginairy/paths.py +++ b/imaginairy/paths.py @@ -1,3 +1,5 @@ +"""Code for defining package root path""" + import os.path PKG_ROOT = os.path.dirname(__file__) diff --git a/imaginairy/prompt_schedules.py b/imaginairy/prompt_schedules.py index 31ba094..a4e85ff 100644 --- a/imaginairy/prompt_schedules.py +++ b/imaginairy/prompt_schedules.py @@ -1,3 +1,5 @@ +"""Functions for managing prompt schedules""" + import csv import re from copy import copy diff --git a/imaginairy/roi_utils.py b/imaginairy/roi_utils.py index 12f1c20..76957b0 100644 --- a/imaginairy/roi_utils.py +++ b/imaginairy/roi_utils.py @@ -1,3 +1,5 @@ +"""Functions for managing image regions of interest""" + import logging logger = logging.getLogger(__name__) diff --git a/imaginairy/safety.py b/imaginairy/safety.py index 5e1d1d1..d0096d7 100644 --- a/imaginairy/safety.py +++ b/imaginairy/safety.py @@ -1,3 +1,5 @@ +"""Classes for image content safety evaluation""" + import logging from functools import lru_cache diff --git a/imaginairy/samplers/base.py b/imaginairy/samplers/base.py index e0203dc..b64d844 100644 --- a/imaginairy/samplers/base.py +++ b/imaginairy/samplers/base.py @@ -1,3 +1,5 @@ +"""Classes and functions for image generation""" + # pylama:ignore=W0613 import logging from abc import ABC diff --git a/imaginairy/samplers/ddim.py b/imaginairy/samplers/ddim.py index f3319e9..ecd28e4 100644 --- a/imaginairy/samplers/ddim.py +++ b/imaginairy/samplers/ddim.py @@ -1,3 +1,5 @@ +"""Classes for diffusion-based image generation""" + # pylama:ignore=W0613 import logging diff --git a/imaginairy/samplers/kdiff.py b/imaginairy/samplers/kdiff.py index 474ba87..715f67f 100644 --- a/imaginairy/samplers/kdiff.py +++ b/imaginairy/samplers/kdiff.py @@ -1,3 +1,5 @@ +"""Classes and functions for image generation with diffusion models""" + # pylama:ignore=W0613 from abc import ABC from typing import Callable diff --git a/imaginairy/samplers/plms.py b/imaginairy/samplers/plms.py index 6406e6e..5e4148e 100644 --- a/imaginairy/samplers/plms.py +++ b/imaginairy/samplers/plms.py @@ -1,3 +1,5 @@ +"""Classes for image generation sampling""" + # pylama:ignore=W0613 import logging diff --git a/imaginairy/schema.py b/imaginairy/schema.py index cecb70a..3406700 100644 --- a/imaginairy/schema.py +++ b/imaginairy/schema.py @@ -1,3 +1,5 @@ +"""Classes for image generation and manipulation""" + # pylint: disable=E0213 import base64 import hashlib diff --git a/imaginairy/surprise_me.py b/imaginairy/surprise_me.py index 71b4a89..74457d5 100644 --- a/imaginairy/surprise_me.py +++ b/imaginairy/surprise_me.py @@ -1,3 +1,5 @@ +"""Functions for generating surprise image edits""" + import logging import os.path diff --git a/imaginairy/train.py b/imaginairy/train.py index 41f4212..4f8f4ac 100644 --- a/imaginairy/train.py +++ b/imaginairy/train.py @@ -1,3 +1,5 @@ +"""Code for training diffusion models""" + import datetime import logging import os diff --git a/imaginairy/training_tools/image_prep.py b/imaginairy/training_tools/image_prep.py index 5a1a331..c508a7a 100644 --- a/imaginairy/training_tools/image_prep.py +++ b/imaginairy/training_tools/image_prep.py @@ -1,3 +1,5 @@ +"""Functions for image preprocessing and generation""" + import logging import os import os.path diff --git a/imaginairy/training_tools/prune_model.py b/imaginairy/training_tools/prune_model.py index 691bd81..f323e74 100644 --- a/imaginairy/training_tools/prune_model.py +++ b/imaginairy/training_tools/prune_model.py @@ -1,3 +1,5 @@ +"""Functions for pruning diffusion models""" + import logging import os diff --git a/imaginairy/training_tools/single_concept.py b/imaginairy/training_tools/single_concept.py index d4849b6..a521fd9 100644 --- a/imaginairy/training_tools/single_concept.py +++ b/imaginairy/training_tools/single_concept.py @@ -1,3 +1,5 @@ +"""Classes for single-concept model finetuning""" + import os import random import re diff --git a/imaginairy/utils/data_distorter.py b/imaginairy/utils/data_distorter.py index a55f6ca..60e70a2 100644 --- a/imaginairy/utils/data_distorter.py +++ b/imaginairy/utils/data_distorter.py @@ -1,3 +1,5 @@ +"""Classes for data distortion testing""" + import contextlib import math import sys diff --git a/imaginairy/utils/debug_info.py b/imaginairy/utils/debug_info.py index be4a0af..43b2050 100644 --- a/imaginairy/utils/debug_info.py +++ b/imaginairy/utils/debug_info.py @@ -1,3 +1,5 @@ +"""Functions for gathering system debug information""" + from functools import lru_cache diff --git a/imaginairy/utils/model_cache.py b/imaginairy/utils/model_cache.py index fcb803f..d8aebd1 100644 --- a/imaginairy/utils/model_cache.py +++ b/imaginairy/utils/model_cache.py @@ -1,3 +1,5 @@ +"""Utilities for managing model caching""" + # pylama: ignore=W0212 import logging from collections import OrderedDict diff --git a/imaginairy/utils/named_resolutions.py b/imaginairy/utils/named_resolutions.py index e10447b..e4c5298 100644 --- a/imaginairy/utils/named_resolutions.py +++ b/imaginairy/utils/named_resolutions.py @@ -1,3 +1,5 @@ +"""Functions for normalizing image resolutions""" + import contextlib _NAMED_RESOLUTIONS = { diff --git a/imaginairy/utils/torch_installer.py b/imaginairy/utils/torch_installer.py index 8519f51..ef06705 100644 --- a/imaginairy/utils/torch_installer.py +++ b/imaginairy/utils/torch_installer.py @@ -1,3 +1,5 @@ +"""Functions for managing PyTorch installation""" + import logging import subprocess diff --git a/imaginairy/vendored/basicsr/arch_util.py b/imaginairy/vendored/basicsr/arch_util.py index 77adb21..6ebbae8 100644 --- a/imaginairy/vendored/basicsr/arch_util.py +++ b/imaginairy/vendored/basicsr/arch_util.py @@ -1,3 +1,5 @@ +"""Functions for neural network weight initialization and layer creation""" + # from https://github.com/XPixelGroup/BasicSR/blob/b0ee3c8414bd39da34f0216cd6bfd8110b85da60/basicsr/archs/arch_util.py # with removals and slightmodifications import logging diff --git a/imaginairy/vendored/basicsr/img_util.py b/imaginairy/vendored/basicsr/img_util.py index 868433c..53249a4 100644 --- a/imaginairy/vendored/basicsr/img_util.py +++ b/imaginairy/vendored/basicsr/img_util.py @@ -1,3 +1,5 @@ +"""Functions for image format conversion""" + # from https://github.com/XPixelGroup/BasicSR/blob/b0ee3c8414bd39da34f0216cd6bfd8110b85da60/basicsr/utils/img_util.py import math import os diff --git a/imaginairy/vendored/basicsr/rrdbnet_arch.py b/imaginairy/vendored/basicsr/rrdbnet_arch.py index 6a2c9ef..765f333 100644 --- a/imaginairy/vendored/basicsr/rrdbnet_arch.py +++ b/imaginairy/vendored/basicsr/rrdbnet_arch.py @@ -1,3 +1,5 @@ +"""Classes for enhanced super-resolution networks""" + # from https://github.com/XPixelGroup/BasicSR/commit/44672b0c6a607fc2a5aaeea70b9f565d1f3f095b # with slight modifications import torch diff --git a/imaginairy/vendored/blip/blip_itm.py b/imaginairy/vendored/blip/blip_itm.py index 9b2b4bd..8e5bfbf 100644 --- a/imaginairy/vendored/blip/blip_itm.py +++ b/imaginairy/vendored/blip/blip_itm.py @@ -1,3 +1,5 @@ +"""Classes for image-text matching""" + import torch import torch.nn.functional as F from models.blip import create_vit, init_tokenizer, load_checkpoint diff --git a/imaginairy/vendored/blip/blip_nlvr.py b/imaginairy/vendored/blip/blip_nlvr.py index 36d7338..1876924 100644 --- a/imaginairy/vendored/blip/blip_nlvr.py +++ b/imaginairy/vendored/blip/blip_nlvr.py @@ -1,3 +1,5 @@ +"""Classes for visual-language reasoning""" + import torch import torch.nn.functional as F from models.blip import create_vit, init_tokenizer, is_url diff --git a/imaginairy/vendored/blip/blip_retrieval.py b/imaginairy/vendored/blip/blip_retrieval.py index 0b6da3f..fa042ed 100644 --- a/imaginairy/vendored/blip/blip_retrieval.py +++ b/imaginairy/vendored/blip/blip_retrieval.py @@ -1,3 +1,5 @@ +"""Classes for image-text retrieval learning""" + import torch import torch.nn.functional as F from models.blip import create_vit, init_tokenizer, load_checkpoint diff --git a/imaginairy/vendored/blip/blip_vqa.py b/imaginairy/vendored/blip/blip_vqa.py index b8816a2..109786c 100644 --- a/imaginairy/vendored/blip/blip_vqa.py +++ b/imaginairy/vendored/blip/blip_vqa.py @@ -1,3 +1,5 @@ +"""Classes for visual question answering""" + import numpy as np import torch import torch.nn.functional as F diff --git a/imaginairy/vendored/blip/nlvr_encoder.py b/imaginairy/vendored/blip/nlvr_encoder.py index d3b085c..382f4b7 100644 --- a/imaginairy/vendored/blip/nlvr_encoder.py +++ b/imaginairy/vendored/blip/nlvr_encoder.py @@ -1,3 +1,5 @@ +"""Classes for transformer-based language modeling""" + import math from typing import Tuple diff --git a/imaginairy/vendored/clip/clip.py b/imaginairy/vendored/clip/clip.py index f1f1a31..197464c 100644 --- a/imaginairy/vendored/clip/clip.py +++ b/imaginairy/vendored/clip/clip.py @@ -1,3 +1,5 @@ +"""Code for interfacing with CLIP models""" + import hashlib import os import urllib diff --git a/imaginairy/vendored/clip/model.py b/imaginairy/vendored/clip/model.py index dbf3e1f..f18c953 100644 --- a/imaginairy/vendored/clip/model.py +++ b/imaginairy/vendored/clip/model.py @@ -1,3 +1,5 @@ +"""Classes for constructing CLIP models""" + from collections import OrderedDict from typing import Tuple, Union diff --git a/imaginairy/vendored/clip/simple_tokenizer.py b/imaginairy/vendored/clip/simple_tokenizer.py index 026b686..f262ee4 100644 --- a/imaginairy/vendored/clip/simple_tokenizer.py +++ b/imaginairy/vendored/clip/simple_tokenizer.py @@ -1,3 +1,5 @@ +"""Tokenizer for byte-pair encoding operations""" + import gzip import html import os diff --git a/imaginairy/vendored/clip/version.py b/imaginairy/vendored/clip/version.py index e69de29..5ff8a7e 100644 --- a/imaginairy/vendored/clip/version.py +++ b/imaginairy/vendored/clip/version.py @@ -0,0 +1 @@ +"""Code for managing CLIP version information""" diff --git a/imaginairy/vendored/codeformer/codeformer_arch.py b/imaginairy/vendored/codeformer/codeformer_arch.py index 75fad6f..3d0451d 100644 --- a/imaginairy/vendored/codeformer/codeformer_arch.py +++ b/imaginairy/vendored/codeformer/codeformer_arch.py @@ -1,3 +1,5 @@ +"""Classes and functions for image feature transformation and synthesis""" + import math from typing import Optional diff --git a/imaginairy/version.py b/imaginairy/version.py index cc00fb6..ae1bea5 100644 --- a/imaginairy/version.py +++ b/imaginairy/version.py @@ -1,3 +1,6 @@ +"""Function for retrieving package version""" + + def get_version(): from importlib.metadata import PackageNotFoundError, version diff --git a/imaginairy/video_sample.py b/imaginairy/video_sample.py index 32c92cb..5169932 100644 --- a/imaginairy/video_sample.py +++ b/imaginairy/video_sample.py @@ -1,3 +1,5 @@ +"""Functions for generating synthetic videos""" + import logging import math import os diff --git a/imaginairy/weight_management/conversion.py b/imaginairy/weight_management/conversion.py index 6f43c64..a7ac60d 100644 --- a/imaginairy/weight_management/conversion.py +++ b/imaginairy/weight_management/conversion.py @@ -1,3 +1,5 @@ +"""Classes for managing model weight conversions""" + import os.path from dataclasses import dataclass from functools import lru_cache diff --git a/imaginairy/weight_management/execution_trace.py b/imaginairy/weight_management/execution_trace.py index 492fba7..fc825b1 100644 --- a/imaginairy/weight_management/execution_trace.py +++ b/imaginairy/weight_management/execution_trace.py @@ -1,3 +1,5 @@ +"""Functions for tracing model execution order""" + import torch from transformers import CLIPTextModelWithProjection diff --git a/imaginairy/weight_management/generate_conversion_maps.py b/imaginairy/weight_management/generate_conversion_maps.py index 544e8bf..2f38ce9 100644 --- a/imaginairy/weight_management/generate_conversion_maps.py +++ b/imaginairy/weight_management/generate_conversion_maps.py @@ -1,3 +1,5 @@ +"""Code for generating model weight conversion maps""" + import itertools import json import os diff --git a/imaginairy/weight_management/generate_weight_info.py b/imaginairy/weight_management/generate_weight_info.py index 723d0b1..748ef5a 100644 --- a/imaginairy/weight_management/generate_weight_info.py +++ b/imaginairy/weight_management/generate_weight_info.py @@ -1,3 +1,5 @@ +"""Functions for managing model weights information""" + import safetensors from imaginairy.model_manager import ( diff --git a/imaginairy/weight_management/pattern_collapse.py b/imaginairy/weight_management/pattern_collapse.py index 77898d8..c99add6 100644 --- a/imaginairy/weight_management/pattern_collapse.py +++ b/imaginairy/weight_management/pattern_collapse.py @@ -1,3 +1,6 @@ +"""Functions for collapsing state_dict key patterns""" + + def find_state_dict_key_patterns(patterns): """Given a list of state_dict keys, collapse similar keys into patterns. diff --git a/imaginairy/weight_management/utils.py b/imaginairy/weight_management/utils.py index 3825889..6b2f2e1 100644 --- a/imaginairy/weight_management/utils.py +++ b/imaginairy/weight_management/utils.py @@ -1,3 +1,5 @@ +"""Utilities for managing AI model weights""" + import os.path _base_dir = os.path.dirname(os.path.realpath(__file__))