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 StableDiffusionPipeline
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 
 from pingpong import PingPong
 from pingpong.pingpong import PPManager
@@ -122,9 +122,15 @@ class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
 TOKEN = os.getenv('HF_TOKEN')
 MODEL_ID = 'meta-llama/Llama-2-70b-chat-hf'
 
-pipe = StableDiffusionPipeline.from_pretrained(
-  "nota-ai/bk-sdm-base", torch_dtype=torch.float16
+pipe = DiffusionPipeline.from_pretrained(
+  "stabilityai/stable-diffusion-2-base", 
+  torch_dtype=torch.float16, 
+  revision="fp16"
 )
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+    pipe.scheduler.config
+)
+pipe = pipe.to("cuda")
 
 STYLES = """
 .left-panel {
@@ -507,7 +513,7 @@ with gr.Blocks(css=STYLES) as demo:
       "This application is designed for you to collaborate with LLM to co-write stories. It is inspired by [Wordcraft project](https://wordcraft-writers-workshop.appspot.com/) from Google's PAIR and Magenta teams. "
       "This application built on [Gradio](https://www.gradio.app), and the underlying text generation is powered by [Hugging Face Inference API](https://huggingface.co/inference-api). The text generation model might"
       "be changed over time, but [meta-llama/Llama-2-70b-chat-hf](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) is selected for now. On the `exporting` tab, you can generate cover art for your story. You can "
-      "design your own prompt to generate the cover art, or you can let the LLaMA2 to generate one for you. Currently, the quantized Stable Diffusion model `nota-ai/bk-sdm-base` by [Nota AI](https://www.nota.ai/) is "
+      "design your own prompt to generate the cover art, or you can let the LLaMA2 to generate one for you. Currently, Stable Diffusion 2.1 model `stabilityai/stable-diffusion-2-base` by [Stability AI](https://stability.ai/) is "
       "used.",
       elem_classes=['center', 'small-big'])
 
@@ -574,7 +580,7 @@ with gr.Blocks(css=STYLES) as demo:
 
       with gr.Tab("Exporting/Sharing"):
         with gr.Column(elem_classes=['group-border']):
-          gr.Markdown("generate cover art with [`nota-ai/bk-sdm-base`](https://huggingface.co/nota-ai/bk-sdm-base) model. "
+          gr.Markdown("generate cover art with [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/stabilityai/stable-diffusion-2-base) model. "
                       "design your own prompt in the textbox below, or just hit 'generate prompt for cover art` button. LLaMA2 "
                       "model will suggest a prompt for you based on your story.")