Update the template

app.py

@@ -3,22 +3,24 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
+#import spaces [uncomment to use ZeroGPU]
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use
 
 if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+    torch_dtype = torch.float16
+else:
+    torch_dtype = torch.float32
+
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+#@spaces.GPU  [uncomment to use ZeroGPU]
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -35,7 +37,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         generator = generator
     ).images[0] 
     
-    return image
+    return image, seed
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -46,21 +48,15 @@ examples = [
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 520px;
+    max-width: 640px;
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
         # Text-to-Image Gradio Template
-        Currently running on {power_device}.
         """)
         
         with gr.Row():
@@ -103,7 +99,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024, #Replace with defaults that work for your model
                 )
                 
                 height = gr.Slider(
@@ -111,7 +107,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024, #Replace with defaults that work for your model
                 )
             
             with gr.Row():
@@ -121,26 +117,26 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,
+                    value=0.0, #Replace with defaults that work for your model
                 )
                 
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
-                    maximum=12,
+                    maximum=50,
                     step=1,
-                    value=2,
+                    value=2, #Replace with defaults that work for your model
                 )
         
         gr.Examples(
             examples = examples,
             inputs = [prompt]
         )
-
-    run_button.click(
+    gr.on(
+        trigger=[run_button.click, prompt.submit],
         fn = infer,
         inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+        outputs = [result, seed]
     )
 
 demo.queue().launch()