Update app.py

app.py

@@ -15,55 +15,275 @@ pipe.to("cuda")
 
 @torch.no_grad()
 def call(
-        pipe, prompt, prompt2, height, width, num_inference_steps, denoising_end,
-        guidance_scale, guidance_scale2, negative_prompt, negative_prompt2, 
-        num_images_per_prompt, eta, generator, latents, prompt_embeds,
-        negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds,
-        output_type, return_dict, callback, callback_steps, cross_attention_kwargs,
-        guidance_rescale, original_size, crops_coords_top_left, target_size,
-        negative_original_size, negative_crops_coords_top_left, negative_target_size):
-    height = height or pipe.default_sample_size * pipe.vae_scale_factor
-    width = width or pipe.default_sample_size * pipe.vae_scale_factor
-    original_size = original_size or (height, width)
-    target_size = target_size or (height, width)
-    pipe.check_inputs(prompt, None, height, width, callback_steps, negative_prompt, None, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds)
-    batch_size = 1 if isinstance(prompt, str) else len(prompt) if isinstance(prompt, list) else prompt_embeds.shape[0]
-    device = pipe._execution_device
-    do_classifier_free_guidance = guidance_scale > 1.0
-    text_encoder_lora_scale = cross_attention_kwargs.get("scale", None) if cross_attention_kwargs else None
-    prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt)
-    prompt2_embeds, negative_prompt2_embeds, pooled_prompt2_embeds, negative_pooled_prompt2_embeds = pipe.encode_prompt(prompt2, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt2)
-    pipe.scheduler.set_timesteps(num_inference_steps, device=device)
-    timesteps = pipe.scheduler.timesteps
-    num_channels_latents = pipe.unet.config.in_channels
-    latents = pipe.prepare_latents(batch_size * num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, latents)
-    extra_step_kwargs = pipe.prepare_extra_step_kwargs(generator, eta)
-    add_text_embeds, add_text2_embeds = pooled_prompt_embeds, pooled_prompt2_embeds
-    add_time_ids = pipe._get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype)
-    add_time2_ids = pipe._get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype=prompt2_embeds.dtype)
-    negative_add_time_ids = pipe._get_add_time_ids(negative_original_size, negative_crops_coords_top_left, negative_target_size, dtype=prompt_embeds.dtype) if negative_original_size and negative_target_size else add_time_ids
-    if do_classifier_free_guidance:
-        prompt_embeds, add_text_embeds, add_time_ids = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0), torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0), torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-        prompt2_embeds, add_text2_embeds, add_time2_ids = torch.cat([negative_prompt2_embeds, prompt2_embeds], dim=0), torch.cat([negative_pooled_prompt2_embeds, add_text2_embeds], dim=0), torch.cat([negative_add_time_ids, add_time2_ids], dim=0)
-    prompt_embeds, add_text_embeds, add_time_ids = prompt_embeds.to(device), add_text_embeds.to(device), add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
-    prompt2_embeds, add_text2_embeds, add_time2_ids = prompt2_embeds.to(device), add_text2_embeds.to(device), add_time2_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
-    num_warmup_steps = max(len(timesteps) - num_inference_steps * pipe.scheduler.order, 0)
-    if denoising_end and isinstance(denoising_end, float) and 0 < denoising_end < 1:
-        discrete_timestep_cutoff = int(round(pipe.scheduler.config.num_train_timesteps - (denoising_end * pipe.scheduler.config.num_train_timesteps)))
-        num_inference_steps = len([ts for ts in timesteps if ts >= discrete_timestep_cutoff])
-        timesteps = timesteps[:num_inference_steps]
-    with pipe.progress_bar(total=num_inference_steps) as progress_bar:
-        for i, t in enumerate(timesteps):
-            if i % 2 == 0:
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
-                noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds, cross_attention_kwargs=cross_attention_kwargs, added_cond_kwargs={"text_embeds": add_text_embeds, "time_ids": add_time_ids})[0]
-                if do_classifier_free_guidance:
-                    noise_pred = noise_pred.chunk(2)[0] + guidance_scale * (noise_pred.chunk(2)[1] - noise_pred.chunk(2)[0])
-            else:
-                latent_model_input2 = torch.cat([latents.flip(2)] * 2) if do_classifier_free_guidance else latents
-                latent_model_input2 = pipe.scheduler.scale_model_input(latent_model_input2, t)
-                noise_pred2 = pipe.unet(latent_model_input2, t)
+        pipe,
+        prompt: Union[str, List[str]] = None,
+        prompt2: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        guidance_scale2: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt2: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+    ):
+        # 0. Default height and width to unet
+        height = height or pipe.default_sample_size * pipe.vae_scale_factor
+        width = width or pipe.default_sample_size * pipe.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        pipe.check_inputs(
+            prompt,
+            None,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            None,
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = pipe._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = pipe.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_embeds=None,
+            negative_prompt_embeds=None,
+            pooled_prompt_embeds=None,
+            negative_pooled_prompt_embeds=None,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        (
+            prompt2_embeds,
+            negative_prompt2_embeds,
+            pooled_prompt2_embeds,
+            negative_pooled_prompt2_embeds,
+        ) = pipe.encode_prompt(
+            prompt=prompt2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt2,
+            prompt_embeds=None,
+            negative_prompt_embeds=None,
+            pooled_prompt_embeds=None,
+            negative_pooled_prompt_embeds=None,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 4. Prepare timesteps
+        pipe.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        timesteps = pipe.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = pipe.unet.config.in_channels
+        latents = pipe.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = pipe.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        add_text2_embeds = pooled_prompt2_embeds
+
+        add_time_ids = pipe._get_add_time_ids(
+            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
+        )
+        add_time2_ids = pipe._get_add_time_ids(
+            original_size, crops_coords_top_left, target_size, dtype=prompt2_embeds.dtype
+        )
+
+        if negative_original_size is not None and negative_target_size is not None:
+            negative_add_time_ids = pipe._get_add_time_ids(
+                negative_original_size,
+                negative_crops_coords_top_left,
+                negative_target_size,
+                dtype=prompt_embeds.dtype,
+            )
+        else:
+            negative_add_time_ids = add_time_ids
+            negative_add_time2_ids = add_time2_ids
+
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+            prompt2_embeds = torch.cat([negative_prompt2_embeds, prompt2_embeds], dim=0)
+            add_text2_embeds = torch.cat([negative_pooled_prompt2_embeds, add_text2_embeds], dim=0)
+            add_time2_ids = torch.cat([negative_add_time2_ids, add_time2_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        prompt2_embeds = prompt2_embeds.to(device)
+        add_text2_embeds = add_text2_embeds.to(device)
+        add_time2_ids = add_time2_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        # 8. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * pipe.scheduler.order, 0)
+
+        # 7.1 Apply denoising_end
+        if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1:
+            discrete_timestep_cutoff = int(
+                round(
+                    pipe.scheduler.config.num_train_timesteps
+                    - (denoising_end * pipe.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+            timesteps = timesteps[:num_inference_steps]
+
+        with pipe.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if i % 2 == 0:
+                  # expand the latents if we are doing classifier free guidance
+                  latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+                  latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
+
+                  # predict the noise residual
+                  added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                  noise_pred = pipe.unet(
+                      latent_model_input,
+                      t,
+                      encoder_hidden_states=prompt_embeds,
+                      cross_attention_kwargs=cross_attention_kwargs,
+                      added_cond_kwargs=added_cond_kwargs,
+                      return_dict=False,
+                  )[0]
+
+                  # perform guidance
+                  if do_classifier_free_guidance:
+                      noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                      noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                else:
+                  # expand the latents if we are doing classifier free guidance
+                  latent_model_input2 = torch.cat([latents.flip(2)] * 2) if do_classifier_free_guidance else latents
+                  latent_model_input2 = pipe.scheduler.scale_model_input(latent_model_input2, t)
+
+                  # predict the noise residual
+                  added_cond2_kwargs = {"text_embeds": add_text2_embeds, "time_ids": add_time2_ids}
+                  noise_pred2 = pipe.unet(
+                      latent_model_input2,
+                      t,
+                      encoder_hidden_states=prompt2_embeds,
+                      cross_attention_kwargs=cross_attention_kwargs,
+                      added_cond_kwargs=added_cond2_kwargs,
+                      return_dict=False,
+                  )[0]
+
+                  # perform guidance
+                  if do_classifier_free_guidance:
+                      noise_pred2_uncond, noise_pred2_text = noise_pred2.chunk(2)
+                      noise_pred2 = noise_pred2_uncond + guidance_scale2 * (noise_pred2_text - noise_pred2_uncond)
+
+                noise_pred = noise_pred if i % 2 == 0 else noise_pred2.flip(2)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = pipe.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipe.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if not output_type == "latent":
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            needs_upcasting = pipe.vae.dtype == torch.float16 and pipe.vae.config.force_upcast
+
+            if needs_upcasting:
+                pipe.upcast_vae()
+                latents = latents.to(next(iter(pipe.vae.post_quant_conv.parameters())).dtype)
+
+            image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
+
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                pipe.vae.to(dtype=torch.float16)
+        else:
+            image = latents
+
+        if not output_type == "latent":
+            # apply watermark if available
+            if pipe.watermark is not None:
+                image = pipe.watermark.apply_watermark(image)
+
+            image = pipe.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        pipe.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusionXLPipelineOutput(images=image)
 
 def simple_call(prompt1, prompt2, guidance_scale1, guidance_scale2, negative_prompt1, negative_prompt2):
     generator = [torch.Generator(device="cuda").manual_seed(5)]