import { Conv2d, GroupNorm, Tensor } from '@jsgrad/jsgrad/base';
import { ClipTextTransformer } from './clip.ts';
import { UNetModel } from './unet.ts';
declare class AttnBlock {
    norm: GroupNorm;
    q: Conv2d;
    k: Conv2d;
    v: Conv2d;
    proj_out: Conv2d;
    constructor(in_channels: number);
    call: (x: Tensor) => Tensor;
}
declare class ResnetBlock {
    norm1: GroupNorm;
    conv1: Conv2d;
    norm2: GroupNorm;
    conv2: Conv2d;
    nin_shortcut?: Conv2d;
    constructor(in_channels: number, out_channels: number);
    call: (x: Tensor) => Tensor;
}
declare class Mid {
    block_1: ResnetBlock;
    attn_1: AttnBlock;
    block_2: ResnetBlock;
    constructor(block_in: number);
    call: (x: Tensor) => Tensor;
}
declare class Decoder {
    conv_in: Conv2d;
    mid: Mid;
    up: {
        block: ResnetBlock[];
        upsample?: {
            conv: Conv2d;
        };
    }[];
    norm_out: GroupNorm;
    conv_out: Conv2d;
    constructor();
    call: (x: Tensor) => Promise<Tensor>;
}
declare class Encoder {
    conv_in: Conv2d;
    down: {
        block: ResnetBlock[];
        downsample?: {
            conv: Conv2d;
        };
    }[];
    mid: Mid;
    norm_out: GroupNorm;
    conv_out: Conv2d;
    constructor();
    call: (x: Tensor) => Tensor;
}
declare class AutoencoderKL {
    encoder: Encoder;
    decoder: Decoder;
    quant_conv: Conv2d;
    post_quant_conv: Conv2d;
    constructor();
    call: (x: Tensor) => Promise<Tensor>;
}
export declare class StableDiffusion {
    alphas_cumprod: Tensor;
    model: {
        diffusion_model: UNetModel;
    };
    first_stage_model: AutoencoderKL;
    cond_stage_model: {
        transformer: {
            text_model: ClipTextTransformer;
        };
    };
    constructor();
    get_x_prev_and_pred_x0: (x: Tensor, e_t: Tensor, a_t: Tensor, a_prev: Tensor) => Tensor[];
    get_model_output: (unconditional_context: Tensor, context: Tensor, latent: Tensor, timestep: Tensor, unconditional_guidance_scale: Tensor) => Tensor;
    decode: (x: Tensor) => Promise<Tensor>;
    call: (unconditional_context: Tensor, context: Tensor, latent: Tensor, timestep: Tensor, alphas: Tensor, alphas_prev: Tensor, guidance: Tensor) => Promise<Tensor>;
}
export {};