update

app.py

@@ -21,26 +21,26 @@ DESCRIPTION += "\n<p style=\"text-align: center\">Unofficial demo for <a href='h
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = False
-#CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
 USE_TORCH_COMPILE = False
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
+#PREVIEW_IMAGES = False
 PREVIEW_IMAGES = True
 
 dtype = torch.bfloat16
-#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-  device = "cuda"
+    device = "cuda"
 elif torch.backends.mps.is_available():
-  device = "mps"
-  dtype = torch.float32
+    device = "mps"
+    dtype = torch.float32
 else:
-  device = "cpu"
+    device = "cpu"
+
 print(f"device={device}")
-#if torch.cuda.is_available():
+
 if device != "cpu":
-    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype, revision="refs/pr/2")#.to(device)
-    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype, revision="refs/pr/44")#.to(device) 
+    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype)#.to(device)
+    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype)#.to(device) 
 
     if ENABLE_CPU_OFFLOAD:
         prior_pipeline.enable_model_cpu_offload()
@@ -57,10 +57,12 @@ if device != "cpu":
         previewer = Previewer()
         previewer_state_dict = torch.load("previewer/previewer_v1_100k.pt", map_location=torch.device('cpu'))["state_dict"]
         previewer.load_state_dict(previewer_state_dict)
-        def callback_prior(i, t, latents):
+        def callback_prior(pipeline, step_index, t, callback_kwargs):
+            latents = callback_kwargs["latents"]
             output = previewer(latents)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
-            return output
+            callback_kwargs["preview_output"] = output
+            return callback_kwargs
         callback_steps = 1
     else:
         previewer = None
@@ -92,14 +94,19 @@ def generate(
     num_images_per_prompt: int = 2,
 #    profile: gr.OAuthProfile | None = None,
 ) -> PIL.Image.Image:
+
     try:
         previewer.eval().requires_grad_(False).to(device).to(dtype)
     except:
         print("")
-    prior_pipeline.to(device)
-    decoder_pipeline.to(device)
     
+    #previewer.eval().requires_grad_(False).to(device).to(dtype)
+#    if device != "cpu":
+#        prior_pipeline.to(device)
+#        decoder_pipeline.to(device)
+#    
     generator = torch.Generator().manual_seed(seed)
+    print("prior_num_inference_steps: ", prior_num_inference_steps)
     prior_output = prior_pipeline(
         prompt=prompt,
         height=height,
@@ -110,15 +117,17 @@ def generate(
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
         generator=generator,
-        callback=callback_prior,
-        callback_steps=callback_steps
+        #callback_on_step_end=callback_prior,
+        #callback_on_step_end_tensor_inputs=['latents']
     )
-
     if PREVIEW_IMAGES:
         for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
-            r = next(prior_output)
-            if isinstance(r, list):
-                yield r[0]
+            try:
+                r = next(prior_output)
+                if isinstance(r, list):
+                    yield r[0]
+            except:
+                print("")
         prior_output = r
 
     decoder_output = decoder_pipeline(
@@ -131,7 +140,6 @@ def generate(
         generator=generator,
         output_type="pil",
     ).images
-
 #    #Save images
 #    for image in decoder_output:
 #        user_history.save_image(
@@ -156,7 +164,7 @@ def generate(
 examples = [
     "An astronaut riding a green horse",
     "A mecha robot in a favela by Tarsila do Amaral",
-    "The sprirt of a Tamagotchi wandering in the city of Los Angeles",
+    "The spirit of a Tamagotchi wandering in the city of Los Angeles",
     "A delicious feijoada ramen dish"
 ]
 
@@ -289,4 +297,4 @@ with gr.Blocks(css="style.css") as demo_with_history:
 #        user_history.render()
 
 if __name__ == "__main__":
-    demo_with_history.queue(max_size=20).launch()
+    demo_with_history.queue(max_size=20).launch()

requirements.txt

@@ -1,6 +1,6 @@
 #git+https://github.com/kashif/diffusers.git@diffusers-yield-callback
 #git+https://github.com/kashif/diffusers.git@a3dc21385b7386beb3dab3a9845962ede6765887
-diffusers
+git+https://github.com/apolinario/diffusers.git@yield-new-pipe
 accelerate
 safetensors
 transformers