Update app and add examples

app.py

@@ -1,57 +1,106 @@
 import gradio as gr
+
 import spaces
-import torch
-import diffusers
 
 from generation import generate_material
 
 
 @spaces.GPU
 def generate(prompts, seed, resolution, refinement):
-  image = generate_material(prompts, seed=seed, resolution=int(resolution), refinement=refinement)
-  return image.basecolor, image.normal, image.height, image.metallic, image.roughness
+    image = generate_material(
+        prompts, seed=seed, resolution=int(resolution), refinement=refinement
+    )
+    return image.basecolor, image.normal, image.height, image.metallic, image.roughness
+    pass
 
 
-def interface_function(prompt_type, text_prompt, image_prompt, seed, resolution, refinement):
+def interface_function(
+    prompt_type, text_prompt, image_prompt, seed, resolution, refinement
+):
     if prompt_type == "Text":
         return generate(text_prompt, seed, resolution, refinement)
     elif prompt_type == "Image":
         return generate(image_prompt, seed, resolution, refinement)
 
+
 def update_visibility(prompt_type):
     if prompt_type == "Text":
         return gr.update(visible=True), gr.update(visible=False)
     elif prompt_type == "Image":
         return gr.update(visible=False), gr.update(visible=True)
 
+
 with gr.Blocks() as demo:
+    title = """
+    <center> 
+
+    <h1> StableMaterials </h1>
+    <b> Try the <a href="https://huggingface.co/gvecchio/StableMaterials">StableMaterials</a> model. Start generating using a text o image prompt, or simply try one of the examples below. </b>
+    <br>
+    <b> NB: </b> The demo currently uses only the base model and does not include the refinement step. The refiner model will be released in a future update.
+
+    </center>
+    """
+    gr.HTML(title)
+
     with gr.Row():
         with gr.Column():
-            prompt_type = gr.Radio(choices=["Text", "Image"], label="Prompt Type", value="Text")
-            text_prompt = gr.Textbox(label="Text Prompt", visible=True, lines=3, placeholder="A brick wall")
+            prompt_type = gr.Radio(
+                choices=["Text", "Image"], label="Prompt Type", value="Text"
+            )
+            text_prompt = gr.Textbox(
+                label="Text Prompt", visible=True, lines=3, placeholder="A brick wall"
+            )
             image_prompt = gr.Image(type="pil", label="Image Prompt", visible=False)
-        
+
         with gr.Column():
             seed = gr.Number(value=-1, label="Seed (-1 for random)")
-            resolution = gr.Dropdown(["512", "1024", "2048"], value="512", label="Resolution", interactive=False)
+            resolution = gr.Dropdown(
+                ["512", "1024", "2048"],
+                value="512",
+                label="Resolution",
+                interactive=False,
+            )
             refinement = gr.Checkbox(label="Refinement", interactive=False)
             generate_button = gr.Button("Generate")
 
-    prompt_type.change(fn=update_visibility, inputs=prompt_type, outputs=[text_prompt, image_prompt])
-    
+    prompt_type.change(
+        fn=update_visibility, inputs=prompt_type, outputs=[text_prompt, image_prompt]
+    )
+
     with gr.Row():
         output_basecolor = gr.Image(label="Base Color", format="png", image_mode="RGB")
         output_normal = gr.Image(label="Normal Map", format="png", image_mode="RGB")
         output_height = gr.Image(label="Height Map", format="png", image_mode="L")
         output_metallic = gr.Image(label="Metallic Map", format="png", image_mode="L")
         output_roughness = gr.Image(label="Roughness Map", format="png", image_mode="L")
-    
+
     generate_button.click(
-        fn=interface_function, 
-        inputs=[prompt_type, text_prompt, image_prompt, seed, resolution, refinement], 
-        outputs=[output_basecolor, output_normal, output_height, output_metallic, output_roughness]
+        fn=interface_function,
+        inputs=[prompt_type, text_prompt, image_prompt, seed, resolution, refinement],
+        outputs=[
+            output_basecolor,
+            output_normal,
+            output_height,
+            output_metallic,
+            output_roughness,
+        ],
+    )
+
+    gr.Markdown("### Examples")
+    gr.Examples(
+        examples=[
+            ["Text", "A brick wall", None],
+            ["Text", "A wooden floor", None],
+            ["Image", None, "./samples/bricks.jpg"],
+            ["Image", None, "./samples/metal.jpg"],
+        ],
+        inputs=[prompt_type, text_prompt, image_prompt],
+    )
+    gr.Markdown(
+        "Try the examples above to see how the Material Generator works. Adjust the parameters and click 'Generate' to create new material maps."
     )
 
+
 if __name__ == "__main__":
     demo.launch()
-