Update app.py

app.py

@@ -11,7 +11,7 @@ import gradio as gr
 from PIL import Image
 from datasets import load_dataset
 from datasets import DownloadMode, VerificationMode
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionXLPipeline
 
 from pingpong import PingPong
 from pingpong.pingpong import PPManager
@@ -122,15 +122,13 @@ class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
 TOKEN = os.getenv('HF_TOKEN')
 MODEL_ID = 'meta-llama/Llama-2-70b-chat-hf'
 
-pipe = DiffusionPipeline.from_pretrained(
-  "stabilityai/stable-diffusion-2-base", 
-  torch_dtype=torch.float16, 
-  revision="fp16"
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", 
+    torch_dtype=torch.float16, 
+    variant="fp16", 
+    use_safetensors=True
 )
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-    pipe.scheduler.config
-)
-pipe = pipe.to("cuda")
+pipe.to("cuda")
 
 STYLES = """
 .margin-auto {
@@ -301,7 +299,7 @@ def gen_art(editor, cover_art_image, gen_cover_art_prompt):
     pipe = pipe.to("cuda")
 
     return [
-        pipe(gen_cover_art_prompt).images[0],
+        pipe(gen_cover_art_prompt, num_inference_steps=25).images[0].resize((512, 512)),
         gen_cover_art_prompt
     ]