Method generate_image

cake-core/src/models/sd/sd.rs:350–545

(
        &mut self,
        args: &ImageGenerationArgs,
        mut callback: F,
    )

#[async_trait]
impl ImageGenerator for SD {
    async fn generate_image<F>(
        &mut self,
        args: &ImageGenerationArgs,
        mut callback: F,
    ) -> Result<(), anyhow::Error>
    where
        F: FnMut(Vec<ImageBuffer<Rgb<u8>, Vec<u8>>>) + Send + 'static,
    {
        use tracing_chrome::ChromeLayerBuilder;
        use tracing_subscriber::prelude::*;

        let ImageGenerationArgs {
            image_prompt,
            uncond_prompt,
            n_steps,
            num_samples,
            bsize,
            tracing,
            guidance_scale,
            img2img,
            img2img_strength,
            image_seed,
            intermediary_images,
            ..
        } = args;

        let sd_version = self.sd_version;

        if !(0. ..=1.).contains(img2img_strength) {
            anyhow::bail!("img2img-strength should be between 0 and 1, got {img2img_strength}")
        }

        let _guard = if *tracing {
            let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
            tracing_subscriber::registry().with(chrome_layer).init();
            Some(guard)
        } else {
            None
        };

        let default_gs = default_guidance_scale(sd_version);
        let guidance_scale = guidance_scale.as_ref().unwrap_or(&default_gs);
        let default_ns = default_n_steps(sd_version);
        let n_steps = n_steps.as_ref().unwrap_or(&default_ns);

        if let Some(seed) = image_seed {
            self.context.device.set_seed(*seed)?;
        }
        let use_guide_scale = guidance_scale > &1.0;

        let mut text_embeddings: Vec<Tensor> = Vec::new();

        let text_embeddings_1 = self
            .text_embeddings(image_prompt, uncond_prompt, use_guide_scale, true)
            .await?;

        text_embeddings.push(text_embeddings_1);