fix diffusers

app.py

@@ -1,11 +1,7 @@
 import numpy as np
 import PIL.Image
 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 fastapi import FastAPI
 import uvicorn
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse, StreamingResponse
@@ -16,7 +12,6 @@ from db import Database
 import uuid
 import logging
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
 from asyncio import Lock
 
 
@@ -40,13 +35,11 @@ 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(
-        "stabilityai/stable-cascade-prior", torch_dtype=dtype
-    )  # .to(device)
+        "stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16
+    ).to(device)
     decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained(
-        "stabilityai/stable-cascade", torch_dtype=dtype
-    )  # .to(device)
-    prior_pipeline.to(device)
-    decoder_pipeline.to(device)
+        "stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.bfloat16
+    ).to(device)
 
     if USE_TORCH_COMPILE:
         prior_pipeline.prior = torch.compile(
@@ -67,16 +60,14 @@ def generate(
     prior_guidance_scale: float = 4.0,
     decoder_num_inference_steps: int = 10,
     decoder_guidance_scale: float = 0.0,
-    num_images_per_prompt: int = 2,
+    num_images_per_prompt: int = 1,
 ) -> PIL.Image.Image:
-
     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,
         negative_prompt=negative_prompt,
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
@@ -133,7 +124,6 @@ async def generate_image(
 
     logging.info(f"Image not found in cache, generating new image")
     async with generate_lock:
-
         pil_image = generate(prompt, negative_prompt, seed)
         img_id = str(uuid.uuid4())
         img_path = IMGS_PATH / f"{img_id}.jpg"