Update app.py

app.py

@@ -1,24 +1,28 @@
 from diffusers import StableDiffusionPipeline
 import gradio as gr
+import torch
+from utils import device
 
-pipe = StableDiffusionPipeline.from_pretrained("nitrosocke/Arcane-Diffusion")
-pipe = pipe.to("cuda")
+pipe = StableDiffusionPipeline.from_pretrained("nitrosocke/Arcane-Diffusion", torch_dtype=torch.float16)
+if torch.cuda.is_available():
+  pipe = pipe.to("cuda")
 
-def inference(prompt, guidance, steps):
-    all_images = [] 
+def inference(prompt, guidance, steps):    
+    all_images = []
     images = pipe([prompt] * 1, num_inference_steps=int(steps), guidance_scale=guidance, width=512, height=512).images
     all_images.extend(images)
     return all_images
 
 with gr.Blocks() as demo:
     with gr.Row():
+        gr.Markdown(f"Running on: {device}")      
         with gr.Column():
             prompt = gr.Textbox(label="prompt")
             guidance = gr.Slider(label="guidance scale", value=7.5, maximum=15)
             steps = gr.Slider(label="steps", value=50, maximum=100, minimum=2)
             run = gr.Button(value="Run")
         with gr.Column():
-            gallery = gr.Gallery(show_label=False)
+            gallery = gr.Gallery()
 
     run.click(inference, inputs=[prompt, guidance, steps], outputs=gallery)
     gr.Examples([
@@ -27,7 +31,7 @@ with gr.Blocks() as demo:
         ["portrait of a beautiful alyx vance half life, volume lighting, concept art, by greg rutkowski!!, colorful, xray melting colors!!, arcane style", 7, 50],
         ["Aloy from Horizon: Zero Dawn, half body portrait, videogame cover art, highly detailed, digital painting, artstation, concept art, smooth, detailed armor, sharp focus, beautiful face, illustration, art by Artgerm and greg rutkowski and alphonse mucha, arcane style", 7, 50],
         ["fantasy portrait painting, digital art, arcane style", 4, 30],
-    ], [prompt, guidance, steps], gallery, inference, cache_examples=False)
+    ], [prompt, guidance, steps], gallery, inference, cache_examples=torch.cuda.is_available())
 
 demo.queue()
 demo.launch()