spandrel.ImageModelDescriptor

class documentation

class ImageModelDescriptor(ModelBase[T], Generic[T]): (source)

Constructor: ImageModelDescriptor(model, state_dict, architecture, purpose, ...)

View In Hierarchy

A model that takes an image as input and returns an image. See __call__ for more information.

Method __call__ Takes a single image tensor as input and returns a single image tensor as output.
Method __init__ Undocumented
Property purpose The purpose of this model.
Instance Variable _call_fn Undocumented
Instance Variable _purpose Undocumented

Inherited from ModelBase:

Method bfloat16 Moves the parameters and buffers of the underlying module to bfloat16 precision.
Method cpu Moves the parameters and buffers of the underlying module to the CPU.
Method cuda Moves the parameters and buffers of the underlying module to the GPU.
Method eval Sets the underlying module in evaluation mode.
Method float Moves the parameters and buffers of the underlying module to single precision (fp32).
Method half Moves the parameters and buffers of the underlying module to half precision (fp16).
Method to Moves and casts the parameters and buffers of the underlying module to the given device and data type.
Method train Sets the underlying module in training mode.
Instance Variable input_channels The number of input image channels of the model. E.g. 3 for RGB, 1 for grayscale.
Instance Variable output_channels The number of output image channels of the model. E.g. 3 for RGB, 1 for grayscale.
Instance Variable scale The output scale of super resolution models. E.g. 4x, 2x, 1x.
Instance Variable size_requirements Size requirements for the input image. E.g. minimum size.
Instance Variable supports_bfloat16 Whether the model supports bfloat16 precision.
Instance Variable supports_half Whether the model supports half precision (fp16).
Instance Variable tags A list of tags for the model, usually describing the size or model parameters. E.g. "64nf" or "large".
Instance Variable tiling Whether the model supports tiling.
Property architecture The architecture of the model.
Property device The device of the underlying module.
Property dtype The data type of the underlying module.
Property model The model itself: a torch.nn.Module with weights loaded in.
Instance Variable _architecture Undocumented
Instance Variable _model Undocumented
@torch.inference_mode()
def __call__(self, image: Tensor) -> Tensor: (source)

Takes a single image tensor as input and returns a single image tensor as output.

The image tensor must be a 4D tensor with shape (1, input_channels, H, W). The data type (float32, float16, bfloat16) and device of the image tensor must be the same as the model. The range of the image tensor must be [0, 1].

The output tensor will be a 4D tensor with shape (1, output_channels, H*scale, W*scale). The data type and device of the output tensor will be the same as the image tensor. The range of the output tensor will be [0, 1].

If the width and height of the image tensor do not satisfy the size_requirements of the model, then the image tensor will be padded to satisfy the requirements. The additional padding will be removed from the output tensor before returning it. If the image already satisfies the requirements, then no padding will be added.

def __init__(self, model: T, state_dict: StateDict, architecture: Architecture[T], purpose: Literal['SR', 'FaceSR', 'Restoration'], tags: list[str], supports_half: bool, supports_bfloat16: bool, scale: int, input_channels: int, output_channels: int, size_requirements: SizeRequirements|None = None, tiling: ModelTiling = ModelTiling.SUPPORTED, call_fn: Callable[[T, Tensor], Tensor]|None = None): (source)
overrides spandrel.ModelBase.__init__

Undocumented

@property
@override
purpose: Literal['SR', 'FaceSR', 'Restoration'] = (source)
overrides spandrel.ModelBase.purpose

The purpose of this model.

_call_fn = (source)

Undocumented

_purpose: Literal['SR', 'FaceSR', 'Restoration'] = (source)

Undocumented