Minor enhancements

app.py

@@ -1,8 +1,4 @@
-import os
-import random
-import uuid
-import json
-
+import os, random, uuid, json
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -10,8 +6,9 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
+DESCRIPTION = None
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
+    DESCRIPTION = "\nRunning on CPU 🥶 This demo may not work on CPU."
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
@@ -21,15 +18,18 @@ ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "sd-community/sdxl-flash",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    add_watermarker=False
+)
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
 if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "sd-community/sdxl-flash",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        add_watermarker=False
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.to("cuda")
+else:
+    pipe.to("cpu")
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
@@ -90,14 +90,15 @@ examples = [
 
 css = '''
 .gradio-container{max-width: 700px !important}
-h1{text-align:center}
+h1{text-align:left}
 footer {
     visibility: hidden
 }
 '''
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("""# SDXL Flash
-        ### First Image processing takes time then images generate faster.""")
+    gr.Markdown(f"""# SDXL Flash
+        ### First Image processing takes time then images generate faster.
+        {DESCRIPTION}""")
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(