UI changes and ZeroGPU optimizations

app.py

@@ -1,55 +1,175 @@
 import gradio as gr
 import spaces
-from omni_zero import OmniZeroSingle
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
-@spaces.GPU(duration=180)
+import sys
+sys.path.insert(0, './diffusers/src')
+
+import torch
+import torch.nn as nn
+
+#Hack for ZeroGPU
+torch.jit.script = lambda f: f
+####
+
+from huggingface_hub import snapshot_download
+from diffusers import DPMSolverMultistepScheduler
+from diffusers.models import ControlNetModel
+
+from transformers import CLIPVisionModelWithProjection
+
+from pipeline import OmniZeroPipeline
+from insightface.app import FaceAnalysis
+from controlnet_aux import ZoeDetector
+from utils import draw_kps, load_and_resize_image, align_images
+
+import cv2
+import numpy as np
+
+base_model="frankjoshua/albedobaseXL_v13"
+
+snapshot_download("okaris/antelopev2", local_dir="./models/antelopev2")
+face_analysis = FaceAnalysis(name='antelopev2', root='./', providers=['CPUExecutionProvider'])
+face_analysis.prepare(ctx_id=0, det_size=(640, 640))
+
+dtype = torch.float16
+
+ip_adapter_plus_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+    "h94/IP-Adapter", 
+    subfolder="models/image_encoder",
+    torch_dtype=dtype,
+).to("cuda")
+
+zoedepthnet_path = "okaris/zoe-depth-controlnet-xl"
+zoedepthnet = ControlNetModel.from_pretrained(zoedepthnet_path,torch_dtype=dtype).to("cuda")
+
+identitiynet_path = "okaris/face-controlnet-xl"
+identitynet = ControlNetModel.from_pretrained(identitiynet_path, torch_dtype=dtype).to("cuda")
+
+zoe_depth_detector = ZoeDetector.from_pretrained("lllyasviel/Annotators").to("cuda")
+
+pipeline = OmniZeroPipeline.from_pretrained(
+    base_model,
+    controlnet=[identitynet, zoedepthnet],
+    torch_dtype=dtype,
+    image_encoder=ip_adapter_plus_image_encoder,
+).to("cuda")
+
+config = pipeline.scheduler.config
+config["timestep_spacing"] = "trailing"
+pipeline.scheduler = DPMSolverMultistepScheduler.from_config(config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++", final_sigmas_type="zero")
+pipeline.load_ip_adapter(["okaris/ip-adapter-instantid", "h94/IP-Adapter", "h94/IP-Adapter"], subfolder=[None, "sdxl_models", "sdxl_models"], weight_name=["ip-adapter-instantid.bin", "ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus_sdxl_vit-h.safetensors"])
+
+def get_largest_face_embedding_and_kps(image, target_image=None):
+    face_info = face_analysis.get(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
+    if len(face_info) == 0:
+        return None, None
+    largest_face = sorted(face_info, key=lambda x: x['bbox'][2] * x['bbox'][3], reverse=True)[0]
+    face_embedding = torch.tensor(largest_face['embedding']).to("cuda")
+    if target_image is None:
+        target_image = image
+    zeros = np.zeros((target_image.size[1], target_image.size[0], 3), dtype=np.uint8)
+    face_kps_image = draw_kps(zeros, largest_face['kps'])
+    return face_embedding, face_kps_image
+
+@spaces.GPU()
 def generate(
-    seed=42,
     prompt="A person",
+    composition_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f",
+    style_image="https://github.com/okaris/omni-zero/assets/1448702/64dc150b-f683-41b1-be23-b6a52c771584",
+    identity_image="https://github.com/okaris/omni-zero/assets/1448702/ba193a3a-f90e-4461-848a-560454531c58",
+    base_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f",
+    seed=42,
     negative_prompt="blurry, out of focus",
     guidance_scale=3.0,
     number_of_images=1,
     number_of_steps=10,
-    base_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f",
     base_image_strength=0.15,
-    composition_image="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f",
     composition_image_strength=1.0,
-    style_image="https://github.com/okaris/omni-zero/assets/1448702/64dc150b-f683-41b1-be23-b6a52c771584",
     style_image_strength=1.0,
-    identity_image="https://github.com/okaris/omni-zero/assets/1448702/ba193a3a-f90e-4461-848a-560454531c58",
     identity_image_strength=1.0,
     depth_image=None,
     depth_image_strength=0.5,
+    progress=gr.Progress(track_tqdm=True)
 ):
+    resolution = 1024
+
+    if base_image is not None:
+        base_image = load_and_resize_image(base_image, resolution, resolution)
+    else:
+        if composition_image is not None:
+            base_image = load_and_resize_image(composition_image, resolution, resolution)
+        else:
+            raise ValueError("You must provide a base image or a composition image")
+
+    if depth_image is None:
+        depth_image = zoe_depth_detector(base_image, detect_resolution=resolution, image_resolution=resolution)
+    else:
+        depth_image = load_and_resize_image(depth_image, resolution, resolution)
+
+    base_image, depth_image = align_images(base_image, depth_image)
+
+    if composition_image is not None:
+        composition_image = load_and_resize_image(composition_image, resolution, resolution)
+    else: 
+        composition_image = base_image
+
+    if style_image is not None:
+        style_image = load_and_resize_image(style_image, resolution, resolution)
+    else:
+        raise ValueError("You must provide a style image")
     
-    omni_zero = OmniZeroSingle(
-        base_model="frankjoshua/albedobaseXL_v13",
-    )
+    if identity_image is not None:
+        identity_image = load_and_resize_image(identity_image, resolution, resolution)
+    else:
+        raise ValueError("You must provide an identity image")
+    
+    face_embedding_identity_image, target_kps = get_largest_face_embedding_and_kps(identity_image, base_image)
+    if face_embedding_identity_image is None:
+        raise ValueError("No face found in the identity image, the image might be cropped too tightly or the face is too small")
+    
+    face_embedding_base_image, face_kps_base_image = get_largest_face_embedding_and_kps(base_image)
+    if face_embedding_base_image is not None:
+        target_kps = face_kps_base_image
 
-    images = omni_zero.generate(
-        seed=seed,
+    pipeline.set_ip_adapter_scale([identity_image_strength,
+        {
+            "down": { "block_2": [0.0, 0.0] },
+            "up": { "block_0": [0.0, style_image_strength, 0.0] }
+        },
+        {
+            "down": { "block_2": [0.0, composition_image_strength] },
+            "up": { "block_0": [0.0, 0.0, 0.0] }
+        }
+    ])
+
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+
+    images = pipeline(
         prompt=prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
-        number_of_images=number_of_images,
-        number_of_steps=number_of_steps,
-        base_image=base_image,
-        base_image_strength=base_image_strength,
-        composition_image=composition_image,
-        composition_image_strength=composition_image_strength,
-        style_image=style_image,
-        style_image_strength=style_image_strength,
-        identity_image=identity_image,
-        identity_image_strength=identity_image_strength,
-        depth_image=depth_image,
-        depth_image_strength=depth_image_strength,
-    )
+        ip_adapter_image=[face_embedding_identity_image, style_image, composition_image],
+        image=base_image,
+        control_image=[target_kps, depth_image],
+        controlnet_conditioning_scale=[identity_image_strength, depth_image_strength],
+        identity_control_indices=[(0,0)],
+        num_inference_steps=number_of_steps, 
+        num_images_per_prompt=number_of_images,
+        strength=(1-base_image_strength),
+        generator=generator,
+        seed=seed,
+    ).images
 
-    # for i, image in enumerate(images):
-    #     image.save(f"oz_output_{i}.jpg")
     return images
 
+#Move the components in the example fields outside so they are available when gr.Examples is instantiated
+
+
 with gr.Blocks() as demo:
+    gr.Markdown("<h1 style='text-align: center'>Omni Zero</h1>")
+    gr.Markdown("<h4 style='text-align: center'>A diffusion pipeline for zero-shot stylized portrait creation [<a href='https://github.com/okaris/omni-zero' target='_blank'>GitHub</a>], [<a href='https://styleof.com/s/remix-yourself' target='_blank'>StyleOf Remix Yourself</a>]</h4>")
     with gr.Row():
         with gr.Column():
             with gr.Row():
@@ -57,38 +177,42 @@ with gr.Blocks() as demo:
             with gr.Row():
                 negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, out of focus")
             with gr.Row():
-                seed = gr.Slider(label="Seed",step=1, minimum=0, maximum=10000000, value=42)
-                number_of_images = gr.Slider(label="Number of Outputs",step=1, minimum=1, maximum=4, value=1)
-            with gr.Row():
-                guidance_scale = gr.Slider(label="Guidance Scale",step=0.1, minimum=0.0, maximum=14.0, value=3.0)
-                number_of_steps = gr.Slider(label="Number of Steps",step=1, minimum=1, maximum=50, value=10)
-            with gr.Row():
-                with gr.Column():
+                with gr.Column(min_width=140):
                     with gr.Row():
-                        base_image = gr.Image(label="Base Image", value="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f")
+                        composition_image = gr.Image(label="Composition")
                     with gr.Row():
-                        base_image_strength = gr.Slider(label="Base Image Strength",step=0.01, minimum=0.0, maximum=1.0, value=0.15)
-                with gr.Column():
+                        composition_image_strength = gr.Slider(label="Strength",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
+            #with gr.Row():
+                with gr.Column(min_width=140):
                     with gr.Row():
-                        composition_image = gr.Image(label="Composition", value="https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f")
+                        style_image = gr.Image(label="Style Image")
                     with gr.Row():
-                        composition_image_strength = gr.Slider(label="Composition Image Strength",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
-            # with gr.Row():
-                with gr.Column():
+                        style_image_strength = gr.Slider(label="Strength",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
+                with gr.Column(min_width=140):
                     with gr.Row():
-                        style_image = gr.Image(label="Style Image", value="https://github.com/okaris/omni-zero/assets/1448702/64dc150b-f683-41b1-be23-b6a52c771584")
+                        identity_image = gr.Image(label="Identity Image")
                     with gr.Row():
-                        style_image_strength = gr.Slider(label="Style Image Strength",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
-                with gr.Column():
-                    with gr.Row():
-                        identity_image = gr.Image(label="Identity Image", value="https://github.com/okaris/omni-zero/assets/1448702/ba193a3a-f90e-4461-848a-560454531c58")
-                    with gr.Row():
-                        identity_image_strength = gr.Slider(label="Identitiy Image Strenght",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
-                # with gr.Column():
+                        identity_image_strength = gr.Slider(label="Strenght",step=0.01, minimum=0.0, maximum=1.0, value=1.0)
+            with gr.Accordion("Advanced options", open=False):
+                with gr.Row():
+                    with gr.Column(min_width=140):
+                            with gr.Row():
+                                base_image = gr.Image(label="Base Image")
+                            with gr.Row():
+                                base_image_strength = gr.Slider(label="Strength",step=0.01, minimum=0.0, maximum=1.0, value=0.15, min_width=120)
+                #    with gr.Column(min_width=140):
                 #     with gr.Row():
                 #         depth_image = gr.Image(label="depth_image", value=None)
                 #     with gr.Row():
                 #         depth_image_strength = gr.Slider(label="depth_image_strength",step=0.01, minimum=0.0, maximum=1.0, value=0.5)
+                            
+                with gr.Row():
+                    seed = gr.Slider(label="Seed",step=1, minimum=0, maximum=10000000, value=42)
+                    number_of_images = gr.Slider(label="Number of Outputs",step=1, minimum=1, maximum=4, value=1)
+                with gr.Row():
+                    guidance_scale = gr.Slider(label="Guidance Scale",step=0.1, minimum=0.0, maximum=14.0, value=3.0)
+                    number_of_steps = gr.Slider(label="Number of Steps",step=1, minimum=1, maximum=50, value=10)
+            
         with gr.Column():
             with gr.Row():
                 out = gr.Gallery(label="Output(s)")
@@ -97,24 +221,30 @@ with gr.Blocks() as demo:
                 submit = gr.Button("Generate")
         
                 submit.click(generate, inputs=[
-                    seed,
                     prompt,
+                    composition_image,
+                    style_image,
+                    identity_image,
+                    base_image,
+                    seed,
                     negative_prompt,
                     guidance_scale,
                     number_of_images,
                     number_of_steps,
-                    base_image,
                     base_image_strength,
-                    composition_image,
                     composition_image_strength,
-                    style_image,
                     style_image_strength,
-                    identity_image,
                     identity_image_strength,
                     ],
                     outputs=[out]
                 )
         # clear.click(lambda: None, None, chatbot, queue=False)
-
+    gr.Examples(
+        examples=[["A person", "https://github.com/okaris/omni-zero/assets/1448702/2ca63443-c7f3-4ba6-95c1-2a341414865f", "https://github.com/okaris/omni-zero/assets/1448702/64dc150b-f683-41b1-be23-b6a52c771584", "https://github.com/okaris/omni-zero/assets/1448702/ba193a3a-f90e-4461-848a-560454531c58"]],
+        inputs=[prompt, composition_image, style_image, identity_image],
+        outputs=[out],
+        fn=generate,
+        cache_examples="lazy",
+    )
 if __name__ == "__main__":
     demo.launch()