Update app.py

update the interface to use Blocks for more customized GUI with examples

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from PIL import Image
 import base64
 import io
+import glob
 import cv2
 import numpy as np
 import torch
@@ -20,24 +21,38 @@ def predict(sketch, description):
     controlnet_id = "lllyasviel/sd-controlnet-scribble"
 
     # Load ControlNet model
-    controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
+    controlnet = ControlNetModel.from_pretrained(controlnet_id)
 
     # Create pipeline with ControlNet model
-    pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=controlnet, torch_dtype=torch.float16)
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=controlnet)
 
     # Use improved scheduler
     pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
     # Enable smart CPU offloading and memory efficient attention
     # pipe.enable_model_cpu_offload()
-    pipe.enable_xformers_memory_efficient_attention()
+    # pipe.enable_xformers_memory_efficient_attention()
 
+    # Move pipeline to GPU
+    pipe = pipe.to("cuda")
 
-    result = pipe(description, image, num_inference_steps=20).images[0]
+    result = pipe(description, image, num_inference_steps=10).images[0]
     
     return result
-# Define sketchpad with custom size and stroke width
-sketchpad = gr.Sketchpad(shape=(1024, 1024), brush_radius=5)
-
-iface = gr.Interface(fn=predict, inputs=[sketchpad, "text"], outputs="image", live=False)
-iface.launch()
+with gr.Blocks() as iface:
+    # Define sketchpad with custom size and stroke width
+    sketchpad = gr.Sketchpad(shape=(400, 300), brush_radius=5, label="Sketchpad- Draw something")
+    txt= gr.Textbox(lines=3, label="Description - Describe your sketch with style")
+    im = gr.Image(label="Output Image", interactive=False)
+    button = gr.Button(value="Submit")
+    button.click(predict, inputs=[sketchpad, txt], outputs=im)
+    flag= gr.CSVLogger()
+    flag.setup([sketchpad, txt, im], "flagged_data_points")
+    button_flag = gr.Button(value="Flag")
+    button_flag.click(lambda *args: flag.flag(args), [sketchpad, txt, im], None, preprocess=False)
+
+    # iface = gr.Interface(fn=predict, inputs=[sketchpad, "text"], outputs=im, live=False, title="Sketch2Image")
+    ## get all the file path from flagged/sketch folder into a list
+    sketch_path = glob.glob("flagged/sketch/*.png")
+    gr.Examples(examples = list(map(lambda x: [x ,"draw in the style of crayon by kids"], sketch_path)), inputs=[sketchpad,txt], outputs=im, fn=predict, cache_examples=True)
+iface.launch()