Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 import spaces
 import numpy as np
 import uuid
-from diffusers import PixArtAlphaPipeline, LCMScheduler, ConsistencyDecoderVAE
+from diffusers import PixArtAlphaPipeline, LCMScheduler
 import torch
 from typing import Tuple
 from datetime import datetime
@@ -78,23 +78,17 @@ STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "(No style)"
 NUM_IMAGES_PER_PROMPT = 1
 
+pipe = PixArtAlphaPipeline.from_pretrained(
+        "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    ).to(device)
+
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     if not negative:
         negative = ""
     return p.replace("{prompt}", positive), n + negative
-
-if torch.cuda.is_available():
-
-    pipe = PixArtAlphaPipeline.from_pretrained(
-        "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-
-    pipe.vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
-
-    pipe.to(device)
         
     # speed-up T5
     pipe.text_encoder.to_bettertransformer()
@@ -108,8 +102,6 @@ def save_image(img):
     img.save(unique_name)
     return unique_name
 
-
-
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)