Update app.py

app.py

@@ -10,7 +10,7 @@ from PIL import Image
 import spaces
 from typing import Tuple
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 
 DESCRIPTION = """ """
 
@@ -34,20 +34,6 @@ MAX_SEED = np.iinfo(np.int32).max
 USE_TORCH_COMPILE = 0
 ENABLE_CPU_OFFLOAD = 0
 
-
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-cascade",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)  
-    pipe.load_lora_weights("dataautogpt3/OpenDalleV1.1", weight_name="OpenDalleV1.1.safetensors", adapter_name="dalle")
-    pipe.set_adapters("dalle")
-    pipe.to("cuda")
-
-
-    
 style_list = [
     {
         "name": "(No style)",
@@ -148,24 +134,36 @@ def generate(
     progress=gr.Progress(track_tqdm=True),
 ):
 
-    
     seed = int(randomize_seed_fn(seed, randomize_seed))
 
     if not use_negative_prompt:
         negative_prompt = ""  # type: ignore
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
 
-    images = pipe(
+    prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
+    decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
+
+    prior.enable_model_cpu_offload()
+    prior_output = prior(
         prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
         height=height,
+        width=width,
+        negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
         num_images_per_prompt=num_images_per_prompt,
-        cross_attention_kwargs={"scale": 0.65},
+        num_inference_steps=num_inference_steps
+    )
+
+    decoder.enable_model_cpu_offload()
+    images = decoder(
+        image_embeddings=prior_output.image_embeddings.to(torch.float16),
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=0.0,
         output_type="pil",
+        num_inference_steps=10
     ).images
+
     image_paths = [save_image(img) for img in images]
     print(image_paths)
     return image_paths, seed
@@ -313,4 +311,4 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
     )
     
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(show_api=False, debug=False)
+    demo.queue(max_size=20).launch(show_api=False, debug=False)