Update app.py

app.py

@@ -30,6 +30,7 @@ pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
     torch_dtype=torch.float16,
 ).to("cuda")
 pipeline.enable_model_cpu_offload()
+pipeline.enable_vae_slicing()
 
 sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,
                                       share_layer_norm=False,
@@ -42,53 +43,64 @@ handler = sa_handler.Handler(pipeline)
 handler.register(sa_args, )
 
 
-# get depth maps
-def get_depth_maps(image):
-    image = load_image(image) #("./example_image/train.png")
-    depth_image1 = pipeline_calls.get_depth_map(image, feature_processor, depth_estimator)
-    #depth_image2 = load_image("./example_image/sun.png").resize((1024, 1024))
-    #mediapy.show_images([depth_image1, depth_image2])
-    return depth_image1 #[depth_image1, depth_image2]
-
 
 
 # run ControlNet depth with StyleAligned
-def style_aligned_controlnet(reference_prompt, target_prompt, image):
-    #reference_prompt = "a poster in flat design style"
-    #target_prompts = [target_prompts] #["a train in flat design style", "the sun in flat design style"]
+def style_aligned_controlnet(ref_style_prompt, depth_map, ref_image, img_generation_prompt):
+    if depth_map == True:
+        image = load_image(ref_image)
+        depth_image = pipeline_calls.get_depth_map(image, feature_processor, depth_estimator)
+    else:
+        depth_image = load_image(ref_image).resize((1024, 1024))
+    #reference_prompt = ref_style_prompt #"a poster in minimalist origami style"
+    #target_prompts = img_generation_prompt #["mona lisa"] #, "gal gadot"]
     controlnet_conditioning_scale = 0.8
-    num_images_per_prompt = 1 # adjust according to VRAM size
-    depth_map = get_depth_maps(image)
+    num_images_per_prompt = 3 # adjust according to VRAM size
     latents = torch.randn(1 + num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)
-    #for deph_map, target_prompt in zip((depth_image1, depth_image2), target_prompts):
     latents[1:] = torch.randn(num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)
-    images = pipeline_calls.controlnet_call(pipeline, [reference_prompt, target_prompt],
-                                            image=depth_map,
+    images = pipeline_calls.controlnet_call(pipeline, [ref_style_prompt, img_generation_prompt],
+                                            image=depth_image,
                                             num_inference_steps=50,
                                             controlnet_conditioning_scale=controlnet_conditioning_scale,
                                             num_images_per_prompt=num_images_per_prompt,
-                                           latents=latents)
-    print(f"images -{images}")
-    return images[0]  
-
-#mediapy.show_images([images[0], deph_map] +  images[1:], titles=["reference", "depth"] + [f'result {i}' for i in range(1, len(images))])
-
+                                            latents=latents)
+    #mediapy.show_images([images[0], depth_image2] +  images[1:], titles=["reference", "depth"] + [f'result {i}' for i in range(1, len(images))])
+    return [images[0], depth_image] +  images[1:], gr.Image(value=images[0], visible=True)
 
 
 with gr.Blocks() as demo:
-    with gr.Row(variant='panel'):
-      with gr.Group():
-        gr.Markdown("### <center>Reference Prompt and Image</center>")
-        ref_prompt = gr.Textbox(label="Enter a Prompt describing the reference image", placeholder='a photo of <object> in <style name> style')
-        depth_map = gr.Image(label="Upload the image to get Depth Map", type='filepath' )
-      with gr.Group():
-        gr.Markdown("### <center>Prompt for generation and generated Image</center>")
-        prompt = gr.Textbox(label="Enter a Prompt", placeholder='a photo of <object> in <style name> style')
-        output = gr.Image(label="Style-Aligned ControlNet",type='pil')
+    
+    with gr.Row():
+      
+      with gr.Column(variant='panel'):
+        ref_style_prompt = gr.Textbox(
+          label='Reference style prompt',
+          info="Enter a Prompt to generate the reference image", placeholder='a poster in <style name> style'
+        )
+        depth_map = gr.Checkbox(label='Depth-map',)
+        ref_style_image = gr.Image(visible=False, label='Reference style image')
+      
+      with gr.Column(variant='panel'): 
+        ref_image = gr.Image(label="Upload the reference image", 
+                             type='filepath' )
+        img_generation_prompt = gr.Textbox(
+            label='ControlNet Prompt',
+            info="Enter a Prompt to generate images using ControlNet and Style-aligned", 
+            )
+    
     btn = gr.Button("Generate", size='sm')
+    gallery = gr.Gallery(label="Style-Aligned ControlNet - Generated images", 
+                           elem_id="gallery",
+                           columns=5, 
+                           rows=1, 
+                           object_fit="contain", 
+                           height="auto",
+                          )
       
-    btn.click(fn=style_aligned_controlnet, inputs=[ref_prompt, prompt, depth_map], outputs=output, api_name="style_aligned_controlnet")
+    btn.click(fn=style_aligned_controlnet, 
+              inputs=[ref_style_prompt, depth_map, ref_image, img_generation_prompt], 
+              outputs=[gallery, ref_style_image], 
+              api_name="style_aligned_controlnet")
 
 
 demo.launch()
-