Update app.py

app.py

@@ -2,6 +2,7 @@ import spaces
 
 import torch
 import torchvision.transforms.functional as TF
+import tomesd
 import numpy as np
 import random
 import os
@@ -21,6 +22,10 @@ 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))
 
@@ -64,6 +69,11 @@ pipe = PhotoMakerStableDiffusionXLAdapterPipeline.from_pretrained(
     variant="fp16",
 ).to(device)
 
+pipe.unet = pipe.unet.to(device=device, dtype=torch_dtype)
+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.load_photomaker_adapter(
     os.path.dirname(photomaker_ckpt),
     subfolder="",
@@ -78,6 +88,11 @@ 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(
@@ -97,82 +112,85 @@ def generate_image(
     adapter_conditioning_factor,
     progress=gr.Progress(track_tqdm=True)
 ):
-    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))
+    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
     
-    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)")
+        # 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)
+    
+        generator = torch.Generator(device=device).manual_seed(seed)
     
-    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)
+        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)
 
 def swap_to_gallery(images):
     return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)