Method load

cake-core/src/models/sd/sd.rs:122–345

(context: &mut Context)

    const MODEL_NAME: &'static str = "stable-diffusion";

    async fn load(context: &mut Context) -> Result<Option<Box<Self>>> {
        let SDArgs {
            tokenizer,
            tokenizer_2,
            sd_version,
            use_f16,
            width,
            height,
            sliced_attention_size,
            clip,
            clip2,
            vae,
            unet,
            use_flash_attention,
            ..
        } = &context.args.sd_args;

        let sd_config = match *sd_version {
            StableDiffusionVersion::V1_5 => {
                StableDiffusionConfig::v1_5(*sliced_attention_size, *height, *width)
            }
            StableDiffusionVersion::V2_1 => {
                StableDiffusionConfig::v2_1(*sliced_attention_size, *height, *width)
            }
            StableDiffusionVersion::Xl => {
                StableDiffusionConfig::sdxl(*sliced_attention_size, *height, *width)
            }
            StableDiffusionVersion::Turbo => {
                StableDiffusionConfig::sdxl_turbo(*sliced_attention_size, *height, *width)
            }
        };

        // Tokenizer
        info!("Loading the Tokenizer...");

        let tokenizer_file = ModelFile::Tokenizer;
        let tokenizer = tokenizer_file.get(
            tokenizer.clone(),
            *sd_version,
            *use_f16,
            context.args.model.clone(),
        )?;
        let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;

        let pad_id = match &sd_config.clip.pad_with {
            Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
            None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
        };

        info!("Tokenizer loaded!");

        // Tokenizer 2

        let mut tokenizer_2_option: Option<Tokenizer> = None;
        let mut pad_id_2: Option<u32> = None;

        if let StableDiffusionVersion::Xl | StableDiffusionVersion::Turbo = sd_version {
            info!("Loading the Tokenizer 2...");