Update app.py

app.py

@@ -2,7 +2,6 @@ import spaces
 
 import torch
 import torchvision.transforms.functional as TF
-import tomesd
 import numpy as np
 import random
 import os
@@ -22,10 +21,6 @@ from aspect_ratio_template import aspect_ratios
 
 # global variable
 base_model_path = 'SG161222/RealVisXL_V5.0'
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
-torch.set_grad_enabled(False)
 face_detector = FaceAnalysis2(providers=['CPUExecutionProvider', 'CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
 face_detector.prepare(ctx_id=0, det_size=(640, 640))
 
@@ -74,6 +69,9 @@ pipe.text_encoder = pipe.text_encoder.to(device=device, dtype=torch_dtype)
 pipe.text_encoder_2 = pipe.text_encoder_2.to(device=device, dtype=torch_dtype)
 pipe.vae = pipe.vae.to(device=device, dtype=torch_dtype)
 
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead")
+pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead")
+
 pipe.load_photomaker_adapter(
     os.path.dirname(photomaker_ckpt),
     subfolder="",
@@ -88,12 +86,9 @@ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.fuse_lora()
 pipe.to(device)
 
+
 pipe.enable_vae_slicing()
 pipe.enable_vae_tiling()
-pipe.enable_xformers_memory_efficient_attention()
-
-torch.cuda.empty_cache()
-
 @spaces.GPU(duration=120)
 def generate_image(
     upload_images, 
@@ -112,85 +107,82 @@ def generate_image(
     adapter_conditioning_factor,
     progress=gr.Progress(track_tqdm=True)
 ):
-    with torch.inference_mode():
-        torch.cuda.empty_cache()
-        if use_doodle:
-            sketch_image = sketch_image["composite"]
-            r, g, b, a = sketch_image.split()
-            sketch_image = a.convert("RGB")
-            sketch_image = TF.to_tensor(sketch_image) > 0.5 # Inversion 
-            sketch_image = TF.to_pil_image(sketch_image.to(torch.float32))
-            adapter_conditioning_scale = adapter_conditioning_scale
-            adapter_conditioning_factor = adapter_conditioning_factor
-        else:
-            adapter_conditioning_scale = 0.
-            adapter_conditioning_factor = 0.
-            sketch_image = None
-    
-        # check the trigger word
-        image_token_id = pipe.tokenizer.convert_tokens_to_ids(pipe.trigger_word)
-        input_ids = pipe.tokenizer.encode(prompt)
-        if image_token_id not in input_ids:
-            raise gr.Error(f"Cannot find the trigger word '{pipe.trigger_word}' in text prompt! Please refer to step 2️⃣")
-    
-        if input_ids.count(image_token_id) > 1:
-            raise gr.Error(f"Cannot use multiple trigger words '{pipe.trigger_word}' in text prompt!")
-    
-        # determine output dimensions by the aspect ratio
-        output_w, output_h = aspect_ratios[aspect_ratio_name]
-        print(f"[Debug] Generate image using aspect ratio [{aspect_ratio_name}] => {output_w} x {output_h}")
-    
-        # apply the style template
-        prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
-    
-        if upload_images is None:
-            raise gr.Error(f"Cannot find any input face image! Please refer to step 1️⃣")
-    
-        input_id_images = []
-        for img in upload_images:
-            input_id_images.append(load_image(img))
-        
-        id_embed_list = []
-    
-        for img in input_id_images:
-            img = np.array(img)
-            img = img[:, :, ::-1]
-            faces = analyze_faces(face_detector, img)
-            if len(faces) > 0:
-                id_embed_list.append(torch.from_numpy((faces[0]['embedding'])))
-    
-        if len(id_embed_list) == 0:
-            raise gr.Error(f"No face detected, please update the input face image(s)")
-        
-        id_embeds = torch.stack(id_embed_list)
+    if use_doodle:
+        sketch_image = sketch_image["composite"]
+        r, g, b, a = sketch_image.split()
+        sketch_image = a.convert("RGB")
+        sketch_image = TF.to_tensor(sketch_image) > 0.5 # Inversion 
+        sketch_image = TF.to_pil_image(sketch_image.to(torch.float32))
+        adapter_conditioning_scale = adapter_conditioning_scale
+        adapter_conditioning_factor = adapter_conditioning_factor
+    else:
+        adapter_conditioning_scale = 0.
+        adapter_conditioning_factor = 0.
+        sketch_image = None
+
+    # check the trigger word
+    image_token_id = pipe.tokenizer.convert_tokens_to_ids(pipe.trigger_word)
+    input_ids = pipe.tokenizer.encode(prompt)
+    if image_token_id not in input_ids:
+        raise gr.Error(f"Cannot find the trigger word '{pipe.trigger_word}' in text prompt! Please refer to step 2️⃣")
+
+    if input_ids.count(image_token_id) > 1:
+        raise gr.Error(f"Cannot use multiple trigger words '{pipe.trigger_word}' in text prompt!")
+
+    # determine output dimensions by the aspect ratio
+    output_w, output_h = aspect_ratios[aspect_ratio_name]
+    print(f"[Debug] Generate image using aspect ratio [{aspect_ratio_name}] => {output_w} x {output_h}")
+
+    # apply the style template
+    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
+
+    if upload_images is None:
+        raise gr.Error(f"Cannot find any input face image! Please refer to step 1️⃣")
+
+    input_id_images = []
+    for img in upload_images:
+        input_id_images.append(load_image(img))
     
-        generator = torch.Generator(device=device).manual_seed(seed)
+    id_embed_list = []
+
+    for img in input_id_images:
+        img = np.array(img)
+        img = img[:, :, ::-1]
+        faces = analyze_faces(face_detector, img)
+        if len(faces) > 0:
+            id_embed_list.append(torch.from_numpy((faces[0]['embedding'])))
+
+    if len(id_embed_list) == 0:
+        raise gr.Error(f"No face detected, please update the input face image(s)")
     
-        print("Start inference...")
-        print(f"[Debug] Seed: {seed}")
-        print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")
-        start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)
-        if start_merge_step > 30:
-            start_merge_step = 30
-        print(start_merge_step)
-        tomesd.apply_patch(pipe, ratio=0.5)
-        images = pipe(
-            prompt=prompt,
-            width=output_w,
-            height=output_h,
-            input_id_images=input_id_images,
-            negative_prompt=negative_prompt,
-            num_images_per_prompt=num_outputs,
-            num_inference_steps=num_steps,
-            start_merge_step=start_merge_step,
-            generator=generator,
-            guidance_scale=guidance_scale,
-            id_embeds=id_embeds,
-            image=sketch_image,
-            adapter_conditioning_scale=adapter_conditioning_scale,
-            adapter_conditioning_factor=adapter_conditioning_factor,
-        ).images
-        return images, gr.update(visible=True)
+    id_embeds = torch.stack(id_embed_list)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    print("Start inference...")
+    print(f"[Debug] Seed: {seed}")
+    print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")
+    start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)
+    if start_merge_step > 30:
+        start_merge_step = 30
+    print(start_merge_step)
+    images = pipe(
+        prompt=prompt,
+        width=output_w,
+        height=output_h,
+        input_id_images=input_id_images,
+        negative_prompt=negative_prompt,
+        num_images_per_prompt=num_outputs,
+        num_inference_steps=num_steps,
+        start_merge_step=start_merge_step,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        id_embeds=id_embeds,
+        image=sketch_image,
+        adapter_conditioning_scale=adapter_conditioning_scale,
+        adapter_conditioning_factor=adapter_conditioning_factor,
+    ).images
+    return images, gr.update(visible=True)
 
 def swap_to_gallery(images):
     return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
@@ -448,4 +440,4 @@ with gr.Blocks(css=css) as demo:
     
     gr.Markdown(article)
     
-demo.launch()
+demo.launch()