Update app.py

app.py

@@ -26,11 +26,11 @@ USE_TORCH_COMPILE = False
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 PREVIEW_IMAGES = True
 
-dtype = torch.float16
+dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    prior_pipeline = StableCascadePriorPipeline.from_pretrained("diffusers/StableCascade-prior", torch_dtype=torch.bfloat16).to("cuda")
-    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("diffusers/StableCascade-decoder",  torch_dtype=torch.bfloat16).to("cuda") 
+    prior_pipeline = StableCascadePriorPipeline.from_pretrained("diffusers/StableCascade-prior", torch_dtype=dtype).to(device)
+    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("diffusers/StableCascade-decoder",  torch_dtype=dtype).to(device) 
 
     if ENABLE_CPU_OFFLOAD:
         prior_pipeline.enable_model_cpu_offload()
@@ -46,6 +46,7 @@ if torch.cuda.is_available():
     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).cpu().numpy())
@@ -81,9 +82,9 @@ def generate(
     num_images_per_prompt: int = 2,
     #profile: gr.OAuthProfile | None = None,
 ) -> PIL.Image.Image:
-    prior_pipeline.to("cuda")
-    decoder_pipeline.to("cuda")
-    previewer.eval().requires_grad_(False).to(device).to(dtype)
+    #prior_pipeline.to(device)
+    #decoder_pipeline.to(device)
+    #previewer.eval().requires_grad_(False).to(device).to(dtype)
     generator = torch.Generator().manual_seed(seed)
     prior_output = prior_pipeline(
         prompt=prompt,