docs: add docstrings

Wrote an openai script and custom prompt to generate them.

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)

imaginairy/animations.py

@@ -1,3 +1,4 @@
+"""Functions for creating animations from images."""
 import os.path
 
 import cv2

imaginairy/api.py

@@ -1,3 +1,5 @@
+"""Functions for generating and processing images"""
+
 import logging
 import os
 import re

imaginairy/api_refiners.py

@@ -1,3 +1,5 @@
+"""Functions for generating refined images"""
+
 import logging
 from typing import List, Optional
 

imaginairy/cli/arg_schedule.py

@@ -1,3 +1,5 @@
+"""Decorator and parsers for command scheduling"""
+
 from typing import Iterable
 
 from imaginairy.utils import frange

imaginairy/cli/colorize.py

@@ -1,3 +1,5 @@
+"""Command for colorizing images with AI"""
+
 import logging
 
 import click

imaginairy/cli/describe.py

@@ -1,3 +1,5 @@
+"""Command for generating image descriptions"""
+
 import click
 
 

imaginairy/cli/edit.py

@@ -1,3 +1,5 @@
+"""Command for AI-powered image editing"""
+
 import click
 
 from imaginairy import config

imaginairy/cli/edit_demo.py

@@ -1,3 +1,5 @@
+"""Command for applying preset image edits"""
+
 import click
 
 

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

imaginairy/cli/main.py

@@ -1,3 +1,5 @@
+"""CLI for AI-powered image generation"""
+
 import logging
 
 import click

imaginairy/cli/run_api.py

@@ -1,3 +1,5 @@
+"""Code for starting an HTTP API server"""
+
 import logging
 
 import click

imaginairy/cli/shared.py

@@ -1,3 +1,5 @@
+"""Context managers and functions for image generation CLI"""
+
 import logging
 import math
 from contextlib import contextmanager

imaginairy/cli/train.py

@@ -1,3 +1,5 @@
+"""CLI commands for model training and image preparation"""
+
 import logging
 
 import click

imaginairy/cli/upscale.py

@@ -1,3 +1,5 @@
+"""Command for upscaling images with AI"""
+
 import logging
 
 import click

imaginairy/cli/videogen.py

@@ -1,3 +1,5 @@
+"""Command for generating AI-powered videos"""
+
 import logging
 
 import click

imaginairy/colorize.py

@@ -1,3 +1,5 @@
+"""Functions for colorizing black and white images"""
+
 import logging
 
 from PIL import Image, ImageEnhance, ImageStat

imaginairy/config.py

@@ -1,3 +1,5 @@
+"""Classes and constants for AI model configuration"""
+
 from dataclasses import dataclass
 from typing import Any, List
 

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

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
 

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

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
 

imaginairy/enhancers/face_restoration_codeformer.py

@@ -1,3 +1,5 @@
+"""Code for enhancing facial images"""
+
 import logging
 from functools import lru_cache
 

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

imaginairy/enhancers/prompt_expansion.py

@@ -1,3 +1,5 @@
+"""Functions for expanding text prompts with phraselists"""
+
 import gzip
 import os.path
 import random

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

imaginairy/enhancers/upscale_riverwing.py

@@ -1,3 +1,5 @@
+"""Classes and functions for image upscaling"""
+
 from functools import lru_cache
 
 import numpy as np

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

imaginairy/http_app/app.py

@@ -1,3 +1,5 @@
+"""FastAPI application for image generation"""
+
 import logging
 import os.path
 import sys

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
 

imaginairy/http_app/stablestudio/routes.py

@@ -1,3 +1,5 @@
+"""API routes for image generation service"""
+
 import uuid
 
 from fastapi import APIRouter

imaginairy/http_app/utils.py

@@ -1,3 +1,5 @@
+"""Functions and class for image generation and encoding"""
+
 import base64
 from io import BytesIO
 

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

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

imaginairy/log_utils.py

@@ -1,3 +1,5 @@
+"""Utilities for image generation logging"""
+
 import logging
 import logging.config
 import re

imaginairy/lr_scheduler.py

@@ -1,3 +1,5 @@
+"""Classes for learning rate scheduling"""
+
 import numpy as np
 
 

imaginairy/model_manager.py

@@ -1,3 +1,5 @@
+"""Classes and functions for managing AI models"""
+
 import logging
 import os
 import re

imaginairy/modules/attention.py

@@ -1,3 +1,5 @@
+"""Classes and functions for attention mechanisms in neural networks"""
+
 import math
 from functools import lru_cache
 

imaginairy/modules/autoencoder.py

@@ -1,3 +1,5 @@
+"""Classes for image autoencoding and manipulation"""
+
 # pylama:ignore=W0613
 import logging
 import math

imaginairy/modules/cldm.py

@@ -1,3 +1,5 @@
+"""Classes for controlled image generation"""
+
 import torch
 from torch import nn
 

imaginairy/modules/clip_embedders.py

@@ -1,3 +1,5 @@
+"""Classes for text and image encoding"""
+
 import kornia
 import torch
 from einops import repeat

imaginairy/modules/diffusion/model.py

@@ -1,3 +1,5 @@
+"""Classes for diffusion-based image generation"""
+
 # pytorch_diffusion + derived encoder decoder
 import gc
 import math

imaginairy/modules/diffusion/openaimodel.py

@@ -1,3 +1,5 @@
+"""Classes for generative diffusion models"""
+
 import math
 from abc import abstractmethod
 from typing import Optional

imaginairy/modules/diffusion/upscaling.py

@@ -1,3 +1,5 @@
+"""Classes for image upscaling diffusion models"""
+
 from functools import partial
 
 import numpy as np

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

imaginairy/modules/distributions.py

@@ -1,3 +1,5 @@
+"""Classes for parametric Gaussian distributions"""
+
 import numpy as np
 import torch
 

imaginairy/modules/ema.py

@@ -1,3 +1,5 @@
+"""Classes for exponential moving average tracking"""
+
 import torch
 from torch import nn
 

imaginairy/modules/encoders.py

@@ -1,3 +1,5 @@
+"""Classes for text encoding transformations"""
+
 import open_clip
 import torch
 from torch import nn

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
 

imaginairy/modules/midas/midas/base_model.py

@@ -1,3 +1,5 @@
+"""Classes for loading model weights"""
+
 import torch
 
 from imaginairy import config

imaginairy/modules/midas/midas/blocks.py

@@ -1,3 +1,5 @@
+"""Functions and classes for neural network construction"""
+
 import torch
 from torch import nn
 

imaginairy/modules/midas/midas/dpt_depth.py

@@ -1,3 +1,5 @@
+"""Classes for depth estimation from images"""
+
 import torch
 from torch import nn
 

imaginairy/modules/midas/midas/transforms.py

@@ -1,3 +1,5 @@
+"""Functions and classes for image preprocessing"""
+
 import math
 
 import cv2

imaginairy/modules/midas/midas/vit.py

@@ -1,3 +1,5 @@
+"""Classes and functions for Vision Transformer processing"""
+
 import math
 import types
 

imaginairy/modules/refiners_sd.py

@@ -1,3 +1,5 @@
+"""Refinement modules for image generation"""
+
 import logging
 import math
 from functools import lru_cache

imaginairy/modules/sgm/attention.py

@@ -1,3 +1,5 @@
+"""Classes for attention-based neural networks"""
+
 import logging
 import math
 from inspect import isfunction

imaginairy/modules/sgm/autoencoder.py

@@ -1,3 +1,5 @@
+"""Classes for autoencoder-based image generation"""
+
 import logging
 import math
 import re

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

imaginairy/modules/sgm/autoencoding/losses/lpips.py

@@ -1,3 +1,5 @@
+"""Classes for latent space perceptual loss"""
+
 import torch
 import torch.nn as nn
 

imaginairy/modules/sgm/autoencoding/lpips/model/model.py

@@ -1,3 +1,5 @@
+"""Classes for image discrimination"""
+
 import functools
 
 import torch.nn as nn

imaginairy/modules/sgm/autoencoding/lpips/util.py

@@ -1,3 +1,5 @@
+"""Utilities for model checkpoint management and normalization layers"""
+
 import hashlib
 import os
 

imaginairy/modules/sgm/autoencoding/lpips/vqperceptual.py

@@ -1,3 +1,5 @@
+"""Functions for discriminator loss calculation"""
+
 import torch
 import torch.nn.functional as F
 

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
 

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

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
 

imaginairy/modules/sgm/diffusion.py

@@ -1,3 +1,5 @@
+"""Classes for stochastic diffusion models"""
+
 import logging
 import math
 from contextlib import contextmanager

imaginairy/modules/sgm/diffusionmodules/denoiser.py

@@ -1,3 +1,5 @@
+"""Classes for image denoising operations"""
+
 from typing import TYPE_CHECKING, Dict, Union
 
 import torch

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
 

imaginairy/modules/sgm/diffusionmodules/denoiser_weighting.py

@@ -1,3 +1,5 @@
+"""Classes for denoiser weighting strategies"""
+
 import torch
 
 

imaginairy/modules/sgm/diffusionmodules/discretizer.py

@@ -1,3 +1,5 @@
+"""Classes for discretizing diffusion steps"""
+
 from abc import abstractmethod
 from functools import partial
 

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

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

imaginairy/modules/sgm/diffusionmodules/loss_weighting.py

@@ -1,3 +1,5 @@
+"""Classes for diffusion loss weighting"""
+
 from abc import ABC, abstractmethod
 
 import torch

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

imaginairy/modules/sgm/diffusionmodules/openaimodel.py

@@ -1,3 +1,5 @@
+"""Classes for generative image modeling"""
+
 import logging
 import math
 from abc import abstractmethod

imaginairy/modules/sgm/diffusionmodules/sampling_utils.py

@@ -1,3 +1,5 @@
+"""Functions for diffusion sampling calculations"""
+
 import torch
 from scipy import integrate
 

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

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

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

imaginairy/modules/sgm/distributions/distributions.py

@@ -1,3 +1,5 @@
+"""Classes for probabilistic distributions handling"""
+
 import numpy as np
 import torch
 

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

imaginairy/modules/sgm/video_attention.py

@@ -1,3 +1,5 @@
+"""Classes for video sequence transformation"""
+
 import logging
 from typing import Optional
 

imaginairy/outpaint.py

@@ -1,3 +1,5 @@
+"""Functions for image outpainting preparation"""
+
 import re
 
 import torch

imaginairy/paths.py

@@ -1,3 +1,5 @@
+"""Code for defining package root path"""
+
 import os.path
 
 PKG_ROOT = os.path.dirname(__file__)

imaginairy/prompt_schedules.py

@@ -1,3 +1,5 @@
+"""Functions for managing prompt schedules"""
+
 import csv
 import re
 from copy import copy

imaginairy/roi_utils.py

@@ -1,3 +1,5 @@
+"""Functions for managing image regions of interest"""
+
 import logging
 
 logger = logging.getLogger(__name__)

imaginairy/safety.py

@@ -1,3 +1,5 @@
+"""Classes for image content safety evaluation"""
+
 import logging
 from functools import lru_cache
 

imaginairy/samplers/base.py

@@ -1,3 +1,5 @@
+"""Classes and functions for image generation"""
+
 # pylama:ignore=W0613
 import logging
 from abc import ABC

imaginairy/samplers/ddim.py

@@ -1,3 +1,5 @@
+"""Classes for diffusion-based image generation"""
+
 # pylama:ignore=W0613
 import logging
 

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

imaginairy/samplers/plms.py

@@ -1,3 +1,5 @@
+"""Classes for image generation sampling"""
+
 # pylama:ignore=W0613
 import logging
 

imaginairy/schema.py

@@ -1,3 +1,5 @@
+"""Classes for image generation and manipulation"""
+
 # pylint: disable=E0213
 import base64
 import hashlib

imaginairy/surprise_me.py

@@ -1,3 +1,5 @@
+"""Functions for generating surprise image edits"""
+
 import logging
 import os.path
 

imaginairy/train.py

@@ -1,3 +1,5 @@
+"""Code for training diffusion models"""
+
 import datetime
 import logging
 import os

imaginairy/training_tools/image_prep.py

@@ -1,3 +1,5 @@
+"""Functions for image preprocessing and generation"""
+
 import logging
 import os
 import os.path

imaginairy/training_tools/prune_model.py

@@ -1,3 +1,5 @@
+"""Functions for pruning diffusion models"""
+
 import logging
 import os
 

imaginairy/training_tools/single_concept.py

@@ -1,3 +1,5 @@
+"""Classes for single-concept model finetuning"""
+
 import os
 import random
 import re

imaginairy/utils/data_distorter.py

@@ -1,3 +1,5 @@
+"""Classes for data distortion testing"""
+
 import contextlib
 import math
 import sys

imaginairy/utils/debug_info.py

@@ -1,3 +1,5 @@
+"""Functions for gathering system debug information"""
+
 from functools import lru_cache
 
 

imaginairy/utils/model_cache.py

@@ -1,3 +1,5 @@
+"""Utilities for managing model caching"""
+
 # pylama: ignore=W0212
 import logging
 from collections import OrderedDict