Add advanced option

app.py

@@ -1,5 +1,7 @@
+from contextlib import nullcontext
 import gradio as gr
 import torch
+from torch import autocast
 from diffusers import DiffusionPipeline
 import streamlit as st
 from transformers import (
@@ -10,6 +12,8 @@ from transformers import (
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 device_dict = {"cuda": 0, "cpu": -1}
+context = autocast if device == "cuda" else nullcontext
+dtype = torch.float16 if device == "cuda" else torch.float32
 
 # Add language detection pipeline
 language_detection_model_ckpt = "papluca/xlm-roberta-base-language-detection"
@@ -30,19 +34,21 @@ pipe = DiffusionPipeline.from_pretrained(
     detection_pipeline=language_detection_pipeline,
     translation_model=trans_model,
     translation_tokenizer=trans_tokenizer,
-    # revision="fp16",
-    # torch_dtype=torch.float16,
+    revision="fp16",
+    torch_dtype=dtype,
 )
 
-pipe.enable_attention_slicing()
 pipe = pipe.to(device)
 
 #torch.backends.cudnn.benchmark = True
 num_samples = 2
 
-def infer(prompt):
-    output = pipe([prompt] * num_samples)
-    return output.images
+def infer(prompt, scale, steps):
+    
+    with context("cuda"):
+        images = pipe(num_samples*[prompt], guidance_scale=scale, num_inference_steps=steps).images
+
+    return images
 
 css = """
         .gradio-container {
@@ -100,7 +106,6 @@ css = """
             border-radius: 14px !important;
         }
         #advanced-options {
-            display: none;
             margin-bottom: 20px;
         }
         .footer {
@@ -167,13 +172,19 @@ block = gr.Blocks(css=css)
 
 examples = [
     [
-        'Una casa en la playa en un atardecer lluvioso'
+        'Una casa en la playa en un atardecer lluvioso',
+        45,
+        7.5,
     ],
     [
-        'Ein Hund, der Orange isst'
+        'Ein Hund, der Orange isst',
+        45,
+        7.5,
     ],
     [
-        "Photo d'un restaurant parisien"
+        "Photo d'un restaurant parisien",
+        45,
+        7.5,
     ],
 ]
 
@@ -216,14 +227,20 @@ with block as demo:
                 )
                
         gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="generated_id").style(
-            grid=[2], height="auto"
+            grid=[1], height="auto"
         )
         
-        ex = gr.Examples(examples=examples, fn=infer, inputs=[text], outputs=gallery, cache_examples=False)
+        with gr.Row(elem_id="advanced-options"):
+            steps = gr.Slider(label="Steps", minimum=5, maximum=50, value=45, step=5)
+            scale = gr.Slider(
+                label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1
+            )
+        
+        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, steps, scale], outputs=gallery, cache_examples=False)
         ex.dataset.headers = [""]
         
-        text.submit(infer, inputs=[text], outputs=gallery)
-        btn.click(infer, inputs=[text], outputs=gallery)
+        text.submit(infer, inputs=[text, steps, scale], outputs=gallery)
+        btn.click(infer, inputs=[text, steps, scale], outputs=gallery)
 
     gr.HTML(
             """