Update app.py

app.py

@@ -8,23 +8,47 @@ keras.mixed_precision.set_global_policy("mixed_float16")
 
 # prepare model
 resolution = 512
-sd_dreambooth_model = models.StableDiffusion(
-    img_width=resolution, img_height=resolution
-)
+sd_dreambooth_model = keras_cv.models.StableDiffusion(
+        img_width=resolution, img_height=resolution, jit_compile=True, 
+    )
 db_diffusion_model = from_pretrained_keras("AmpleBasis/seymour-cat")
 sd_dreambooth_model._diffusion_model = db_diffusion_model
 
-# generate images
-def infer(prompt):
-    generated_images = sd_dreambooth_model.text_to_image(
-        prompt, batch_size=2
+def generate_images(prompt: str, negative_prompt:str, num_imgs_to_gen: int, num_steps: int, ugs: int):
+    generated_img = sd_dreambooth_model.text_to_image(
+        prompt, 
+        negative_prompt=negative_prompt,
+        batch_size=num_imgs_to_gen,
+        num_steps=num_steps,
+        unconditional_guidance_scale=ugs,
     )
-    return generated_images 
+   
+    return generated_img
+    
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # Seymour Diffusion
+    This is a Keras Dreambooth model fine-tuned to images of Seymour, a cat.
+    The model, part of the [Keras Dreambooth Sprint](https://github.com/huggingface/community-events/tree/main/keras-dreambooth-sprint) trained by Pedro Pacheco, can be found in [AmpleBasis/seymour-cat](https://huggingface.co/AmpleBasis/seymour-cat).
+    
+    The model should be used with a prompt containing `symr cat`. A typical prompt for this model is `photo of symr cat`.
     
-output = gr.Gallery(label="Outputs").style(grid=(1,2))
+    """)
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(lines=1, value="photo of symr cat", label="Prompt")
+            negative_prompt = gr.Textbox(lines=1, value="deformed,blurry,lowres", label="Negative Prompt")
+            samples = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of Images")
+            num_steps = gr.Slider(label="Steps",value=40)
+            ugs = gr.Slider(value=7, minimum=5, maximum=25, step=1, label="Unconditional Guidance Scale")
+            run = gr.Button(value="Generate")
+        with gr.Column():
+            gallery = gr.Gallery(label="Outputs").style(grid=(1,2))
+
+    run.click(generate_images, inputs=[prompt,negative_prompt, samples, num_steps, ugs], outputs=gallery)
+    
+    gr.Examples([["photo of symr cat", 1, 40, 7],
+
+                [prompt,negative_prompt, samples,num_steps, ugs], gallery, generate_images)
 
-# customize interface
-title = "Dreambooth Demo on Seymour the Cat"
-description = "This is a dreambooth model fine-tuned on images of Seymour the Cat. To try it, input the concept with {symr cat}."
-examples=[["photo of symr cat"]]
-gr.Interface(infer, inputs=["text"], outputs=[output], title=title, description=description, examples=examples).queue().launch()
+demo.launch(debug=True)