Upload layer_diff_dataset/test_inp_sd_1.py with huggingface_hub

layer_diff_dataset/test_inp_sd_1.py

@@ -0,0 +1,72 @@
+from diffusers import AutoPipelineForInpainting
+from diffusers.utils import load_image
+import torch
+import os 
+from tqdm import tqdm
+import cv2
+from PIL import Image
+
+pipe = AutoPipelineForInpainting.from_pretrained("../alpha_work/diffusers/stable-diffusion-xl-1.0-inpainting_", torch_dtype=torch.float16, variant="fp16").to("cuda")
+# print('pipe',pipe)
+# StableDiffusionXLInpaintPipeline
+
+preprocessed_root_path = '../data/video_dataset/YoutubeVOS/train/ra_preprocessed'
+# folder_path_0 = '../codes/Inpaint-Anything/results/0b6f9105fc'
+mask_root_path = '../data/video_dataset/YoutubeVOS/train/mask'
+output_root_path = '../data/video_dataset/YoutubeVOS/train/inp_preprocess_sd_0.9_base'
+os.makedirs(output_root_path,exist_ok=True)
+# img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+# mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+
+
+prompt = "hazy background with nothing on"
+generator = torch.Generator(device="cuda").manual_seed(22)
+# base_image = Image.open(base_image_path).resize((1024, 1024))
+
+vid_list = os.listdir(preprocessed_root_path)
+pbar = tqdm(enumerate(vid_list),total=len(vid_list))
+for i, vid_name in pbar:
+    # if idx < 800:
+    #         continue
+    if i>=800:
+        break
+    if i%100==0:
+        print('i: ',i)
+    img_folder = os.path.join(preprocessed_root_path, vid_name)
+    img_list = os.listdir(img_folder)
+    img_list = [i for i in img_list if i.endswith('.png')]
+    img_list.sort()
+    # print(len(img_list))
+    mask_folder = os.path.join(mask_root_path, vid_name)
+    output_folder = os.path.join(output_root_path, vid_name)
+    os.makedirs(output_folder,exist_ok=True)
+
+    # if os.path.exists(os.path.join(folder_path_,image_name)):
+    #     continue
+    for i,image_name in enumerate(img_list):
+        image_path = os.path.join(img_folder,image_name)
+        mask_path = os.path.join(mask_folder,image_name)
+        image = Image.open(image_path).resize((1024, 1024))
+        mask_image = Image.open(mask_path).resize((1024, 1024))
+        # image = cv2.resize(cv2.imread(image_path),(1024,1024))
+        # mask_image = cv2.resize(cv2.imread(mask_path,cv2.IMREAD_GRAYSCALE),(1024,1024))
+        # image = load_image(img_url).resize((1024, 1024))
+        # mask_image = load_image(mask_url).resize((1024, 1024))
+        if i==0:
+            base_image = image
+            # strength = 0.99
+        image_out = pipe(
+        prompt=prompt,
+        image=image,
+        base_image=base_image,
+        mask_image=mask_image,
+        guidance_scale=8.0,
+        num_inference_steps=20,  # steps between 15 and 30 work well for us
+        strength=0.9,  # make sure to use `strength` below 1.0
+        generator=generator,
+        ).images[0]
+        image_out.save(os.path.join(output_folder,image_name))
+        if i==0:
+            base_image = image_out
+            # strength = 0.8
+    # exit(0)