moved to main pipe

app.py

@@ -7,7 +7,6 @@ import torch
 from typing import List
 from diffusers.utils import numpy_to_pil
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
-from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
 from previewer.modules import Previewer
 import os
 import datetime
@@ -61,19 +60,6 @@ if torch.cuda.is_available():
         prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
         decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="max-autotune", fullgraph=True)
     
-    if PREVIEW_IMAGES:
-        previewer = Previewer()
-        previewer.load_state_dict(torch.load("previewer/previewer_v1_100k.pt")["state_dict"])
-        previewer.eval().requires_grad_(False).to(device).to(dtype)
-        def callback_prior(i, t, latents):
-            output = previewer(latents)
-            output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
-            return output
-        callback_steps = 1
-    else:
-        previewer = None
-        callback_prior = None
-        callback_steps = None
 else:
     prior_pipeline = None
     decoder_pipeline = None
@@ -96,40 +82,33 @@ def generate(
     decoder_guidance_scale: float = 0.0,
     batch_size_per_prompt: int = 2,
     number_of_images_per_prompt: int = 1,  # New parameter
-) -> List[PIL.Image.Image]:
+):
     images = []  # Initialize an empty list to collect generated images
     original_seed = seed  # Store the original seed value
     for i in range(number_of_images_per_prompt):
         if i > 0:  # Update seed for subsequent iterations
             seed = random.randint(0, MAX_SEED)
         generator = torch.Generator().manual_seed(seed)
+
         prior_output = prior_pipeline(
             prompt=prompt,
             height=height,
             width=width,
-            num_inference_steps=prior_num_inference_steps,
-            timesteps=DEFAULT_STAGE_C_TIMESTEPS,
+            generator=generator,
             negative_prompt=negative_prompt,
             guidance_scale=prior_guidance_scale,
             num_images_per_prompt=batch_size_per_prompt,
-            generator=generator,
-            callback=callback_prior,
-            callback_steps=callback_steps
+            num_inference_steps=prior_num_inference_steps
         )
 
-        if PREVIEW_IMAGES:
-            for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
-                r = next(prior_output)
-            prior_output = r
-
         decoder_output = decoder_pipeline(
             image_embeddings=prior_output.image_embeddings,
             prompt=prompt,
-            num_inference_steps=            decoder_num_inference_steps,
-            guidance_scale=decoder_guidance_scale,
             negative_prompt=negative_prompt,
-            generator=generator,
+            guidance_scale=decoder_guidance_scale,
             output_type="pil",
+            generator=generator,
+            num_inference_steps=decoder_num_inference_steps
         ).images
 
         # Append generated images to the images list
@@ -233,7 +212,7 @@ with gr.Blocks() as app:
                         minimum=1,
                         maximum=100,
                         step=1,
-                        value=20,
+                        value=30,
                     )
                 with gr.Column():
                     decoder_num_inference_steps = gr.Slider(

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/kashif/diffusers.git@diffusers-yield-callback
+git+https://github.com/kashif/diffusers.git@wuerstchen-v3
 accelerate
 safetensors
 transformers