update

diffusion_webui/stable_diffusion/inpaint_app.py

@@ -8,7 +8,6 @@ stable_inpiant_model_list = [
 ]
 
 stable_prompt_list = ["a photo of a man.", "a photo of a girl."]
-
 stable_negative_prompt_list = ["bad, ugly", "deformed"]
 
 

diffusion_webui/stable_diffusion/keras_txt2img.py

@@ -1,51 +1,55 @@
 import gradio as gr
-import tensorflow as tf
 from huggingface_hub import from_pretrained_keras
 from keras_cv import models
 from tensorflow import keras
 
 keras_model_list = [
     "keras-dreambooth/keras_diffusion_lowpoly_world",
-    "keras-dreambooth/pink-floyd-division-bell",
-    "keras-dreambooth/dreambooth_diffusion_model",
+    "keras-dreambooth/keras-diffusion-traditional-furniture",
 ]
 
 stable_prompt_list = [
-    "a photo of lowpoly_world",
-    "Flower vase inspired by pink floyd division bell",
+    "photo of lowpoly_world",
+    "photo of traditional_furniture",
 ]
 
 stable_negative_prompt_list = ["bad, ugly", "deformed"]
 
+keras.mixed_precision.set_global_policy("mixed_float16")
+dreambooth_model = models.StableDiffusion(
+    img_width=512,
+    img_height=512,
+    jit_compile=True,
+)
+
 
 def keras_stable_diffusion(
     model_path: str,
     prompt: str,
     negative_prompt: str,
-    guidance_scale: int,
-    num_inference_step: int,
-    height: int,
-    width: int,
+    num_imgs_to_gen: int,
+    num_steps: int,
 ):
-
-    with tf.device("/GPU:0"):
-        keras.mixed_precision.set_global_policy("mixed_float16")
-
-        sd_dreambooth_model = models.StableDiffusion(
-            img_width=height, img_height=width
-        )
-
-        db_diffusion_model = from_pretrained_keras(model_path)
-        sd_dreambooth_model._diffusion_model = db_diffusion_model
-
-        generated_images = sd_dreambooth_model.text_to_image(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_steps=num_inference_step,
-            unconditional_guidance_scale=guidance_scale,
-        )
-
-    return generated_images
+    """
+    This function is used to generate images using our fine-tuned keras dreambooth stable diffusion model.
+    Args:
+        prompt (str): The text input given by the user based on which images will be generated.
+        num_imgs_to_gen (int): The number of images to be generated using given prompt.
+        num_steps (int): The number of denoising steps
+    Returns:
+        generated_img (List): List of images that were generated using the model
+    """
+    loaded_diffusion_model = from_pretrained_keras(model_path)
+    dreambooth_model._diffusion_model = loaded_diffusion_model
+
+    generated_img = dreambooth_model.text_to_image(
+        prompt,
+        negative_prompt=negative_prompt,
+        batch_size=num_imgs_to_gen,
+        num_steps=num_steps,
+    )
+
+    return generated_img
 
 
 def keras_stable_diffusion_app():
@@ -68,43 +72,26 @@ def keras_stable_diffusion_app():
                     label="Negative Prompt",
                 )
 
-                with gr.Accordion("Advanced Options", open=False):
-                    keras_text2image_guidance_scale = gr.Slider(
-                        minimum=0.1,
-                        maximum=15,
-                        step=0.1,
-                        value=7.5,
-                        label="Guidance Scale",
-                    )
-
-                    keras_text2image_num_inference_step = gr.Slider(
-                        minimum=1,
-                        maximum=100,
-                        step=1,
-                        value=50,
-                        label="Num Inference Step",
-                    )
-
-                    keras_text2image_height = gr.Slider(
-                        minimum=128,
-                        maximum=1280,
-                        step=32,
-                        value=512,
-                        label="Image Height",
-                    )
-
-                    keras_text2image_width = gr.Slider(
-                        minimum=128,
-                        maximum=1280,
-                        step=32,
-                        value=512,
-                        label="Image Height",
-                    )
+                keras_text2image_guidance_scale = gr.Slider(
+                    minimum=0.1,
+                    maximum=15,
+                    step=0.1,
+                    value=7.5,
+                    label="Guidance Scale",
+                )
+
+                keras_text2image_num_inference_step = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=50,
+                    label="Num Inference Step",
+                )
 
                 keras_text2image_predict = gr.Button(value="Generator")
 
             with gr.Column():
-                output_image = gr.Gallery(label="Output")
+                output_image = gr.Gallery(label="Outputs").style(grid=(1, 2))
 
         gr.Examples(
             fn=keras_stable_diffusion,
@@ -114,8 +101,6 @@ def keras_stable_diffusion_app():
                 keras_text2image_negative_prompt,
                 keras_text2image_guidance_scale,
                 keras_text2image_num_inference_step,
-                keras_text2image_height,
-                keras_text2image_width,
             ],
             outputs=[output_image],
             examples=[
@@ -141,8 +126,6 @@ def keras_stable_diffusion_app():
                 keras_text2image_negative_prompt,
                 keras_text2image_guidance_scale,
                 keras_text2image_num_inference_step,
-                keras_text2image_height,
-                keras_text2image_width,
             ],
             outputs=output_image,
         )