chores: rebase commits

threestudio/scripts/dreamcraft3d_dreambooth.py

@@ -0,0 +1,53 @@
+import argparse
+import os
+from subprocess import run, CalledProcessError
+
+import cv2
+import glob
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from tqdm import tqdm
+from torchvision.utils import save_image
+
+from threestudio.scripts.generate_mv_datasets import generate_mv_dataset
+from threestudio.utils.config import load_config
+import threestudio
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="path to config file")
+    parser.add_argument("--action", default="both", help="action to perform", choices=["gen_data", "dreambooth", "both""])
+    args, extras = parser.parse_known_args()
+    return args, extras
+
+
+def main(args, extras):
+    cfg = load_config(args.config, cli_args=extras, n_gpus=1)
+
+    if args.action == "gen_data" or args.action == "both":
+        # Generate multi-view dataset
+        generate_mv_dataset(cfg)
+
+    if args.action == "dreambooth" or args.action == "both":
+        # Run DreamBooth.
+        command = f'accelerate launch threestudio/scripts/train_dreambooth.py \
+                    --pretrained_model_name_or_path="{cfg.custom_import.dreambooth.model_name}" \
+                    --instance_data_dir="{cfg.custom_import.dreambooth.instance_dir}" \
+                    --output_dir="{cfg.custom_import.dreambooth.output_dir}"\
+                    --instance_prompt="{cfg.custom_import.dreambooth.prompt_dreambooth}" \
+                    --resolution=512 \
+                    --train_batch_size=2 \
+                    --gradient_accumulation_steps=1 \
+                    --learning_rate=1e-6 \
+                    --lr_scheduler="constant" \
+                    --lr_warmup_steps=0 \
+                    --max_train_steps=1000'
+
+        os.system(command)
+
+
+if __name__ == "__main__":
+    args, extras = parse_args()
+    main(args, extras)

threestudio/scripts/generate_images_if.py

@@ -0,0 +1,92 @@
+from diffusers import DiffusionPipeline
+from diffusers.utils import pt_to_pil
+import torch
+
+import os
+import glob
+import json
+import argparse
+import numpy as np
+from tqdm import tqdm
+
+
+SAVE_FOLDER = "./load/images_dreamfusion"
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rank", default=0, type=int, help="# of GPU")
+    parser.add_argument("--prompt",required=True, type=str)
+
+    args = parser.parse_args()
+
+    # stage 1
+    stage_1 = DiffusionPipeline.from_pretrained(
+        "DeepFloyd/IF-I-XL-v1.0",
+        variant="fp16",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_1.enable_model_cpu_offload()
+
+    # stage 2
+    stage_2 = DiffusionPipeline.from_pretrained(
+        "DeepFloyd/IF-II-L-v1.0",
+        text_encoder=None,
+        variant="fp16",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    # stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_2.enable_model_cpu_offload()
+
+    # stage 3
+    # safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker, "watermarker": stage_1.watermarker}
+    safety_modules = None
+    stage_3 = DiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-x4-upscaler",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_3.enable_model_cpu_offload()
+
+    # # load prompt library
+    # with open(os.path.join("load/prompt_library.json"), "r") as f:
+    #     prompt_library = json.load(f)
+
+    # n_prompts = len(prompt_library["dreamfusion"]) 
+    # n_prompts_per_rank = int(np.ceil(n_prompts / 8))
+
+    # for prompt in tqdm(prompt_library["dreamfusion"][args.rank * n_prompts_per_rank : (args.rank + 1) * n_prompts_per_rank]):
+    
+    prompt = args.prompt
+    print("Prompt:", prompt)
+
+    save_folder = os.path.join(SAVE_FOLDER, prompt)
+    os.makedirs(save_folder, exist_ok=True)
+
+    # if len(glob.glob(f"{save_folder}/*.png")) >= 30:
+    #     continue
+
+    # enhance prompt
+    prompt = prompt + ", 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3, hyperrealistic, intricate details, ultra-realistic, award-winning"
+
+    prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
+    for _ in tqdm(range(30)):
+        seed = np.random.randint(low=0, high=10000000, size=1)[0]
+        generator = torch.manual_seed(seed)
+
+        ### Stage 1
+        image = stage_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt").images
+        # pt_to_pil(image)[0].save("./if_stage_I.png")
+
+        ### Stage 2
+        image = stage_2(
+            image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt"
+        ).images
+        # pt_to_pil(image)[0].save("./if_stage_II.png")
+
+        ### Stage 3
+        image = stage_3(prompt=prompt, image=(image.float() * 0.5 + 0.5), generator=generator, noise_level=100).images
+        image[0].save(f"{save_folder}/img_{seed:08d}.png")

threestudio/scripts/generate_images_if_prompt_library.py

@@ -0,0 +1,90 @@
+from diffusers import DiffusionPipeline
+from diffusers.utils import pt_to_pil
+import torch
+
+import os
+import glob
+import json
+import argparse
+import numpy as np
+from tqdm import tqdm
+
+
+SAVE_FOLDER = "./load/images_dreamfusion"
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rank", default=0, type=int, help="# of GPU")
+
+    args = parser.parse_args()
+
+    # stage 1
+    stage_1 = DiffusionPipeline.from_pretrained(
+        "DeepFloyd/IF-I-XL-v1.0",
+        variant="fp16",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_1.enable_model_cpu_offload()
+
+    # stage 2
+    stage_2 = DiffusionPipeline.from_pretrained(
+        "DeepFloyd/IF-II-L-v1.0",
+        text_encoder=None,
+        variant="fp16",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    # stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_2.enable_model_cpu_offload()
+
+    # stage 3
+    # safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker, "watermarker": stage_1.watermarker}
+    safety_modules = None
+    stage_3 = DiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-x4-upscaler",
+        torch_dtype=torch.float16,
+        local_files_only=True
+    )
+    stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+    stage_3.enable_model_cpu_offload()
+
+    # load prompt library
+    with open(os.path.join("load/prompt_library.json"), "r") as f:
+        prompt_library = json.load(f)
+
+    n_prompts = len(prompt_library["dreamfusion"]) 
+    n_prompts_per_rank = int(np.ceil(n_prompts / 8))
+
+    for prompt in tqdm(prompt_library["dreamfusion"][args.rank * n_prompts_per_rank : (args.rank + 1) * n_prompts_per_rank]):
+
+        print("Prompt:", prompt)
+
+        save_folder = os.path.join(SAVE_FOLDER, prompt)
+        os.makedirs(save_folder, exist_ok=True)
+
+        if len(glob.glob(f"{save_folder}/*.png")) >= 30:
+            continue
+
+        # enhance prompt
+        prompt = prompt + ", 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3, hyperrealistic, intricate details, ultra-realistic, award-winning"
+
+        prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
+        for _ in tqdm(range(30)):
+            seed = np.random.randint(low=0, high=10000000, size=1)[0]
+            generator = torch.manual_seed(seed)
+
+            ### Stage 1
+            image = stage_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt").images
+            # pt_to_pil(image)[0].save("./if_stage_I.png")
+
+            ### Stage 2
+            image = stage_2(
+                image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt"
+            ).images
+            # pt_to_pil(image)[0].save("./if_stage_II.png")
+
+            ### Stage 3
+            image = stage_3(prompt=prompt, image=(image.float() * 0.5 + 0.5), generator=generator, noise_level=100).images
+            image[0].save(f"{save_folder}/img_{seed:08d}.png")

threestudio/scripts/generate_mv_datasets.py

@@ -0,0 +1,95 @@
+import os
+import cv2
+import glob
+import torch
+import argparse
+import numpy as np
+from tqdm import tqdm
+import pytorch_lightning as pl
+from torchvision.utils import save_image
+from subprocess import run, CalledProcessError
+from threestudio.utils.config import load_config
+import threestudio
+
+# Constants
+AZIMUTH_FACTOR = 360
+IMAGE_SIZE = (512, 512)
+
+
+def copy_file(source, destination):
+    try:
+        command = ['cp', source, destination]
+        result = run(command, capture_output=True, text=True)
+        result.check_returncode()
+    except CalledProcessError as e:
+        print(f'Error: {e.output}')
+
+
+def prepare_images(cfg):
+    rgb_list = sorted(glob.glob(os.path.join(cfg.data.render_image_path, "*.png")))
+    rgb_list.sort(key=lambda file: int(os.path.splitext(os.path.basename(file))[0]))
+    n_rgbs = len(rgb_list)
+    n_samples = cfg.data.n_samples
+
+    os.makedirs(cfg.data.save_path, exist_ok=True)
+
+    copy_file(cfg.data.ref_image_path, f"{cfg.data.save_path}/ref_0.0.png")
+
+    sampled_indices = np.linspace(0, len(rgb_list)-1, n_samples, dtype=int)
+    rgb_samples = [rgb_list[index] for index in sampled_indices]
+
+    return rgb_samples
+
+
+def process_images(rgb_samples, cfg, guidance, prompt_utils):
+    n_rgbs = 120
+    for rgb_name in tqdm(rgb_samples):
+        rgb_idx = int(os.path.basename(rgb_name).split(".")[0])
+        rgb = cv2.imread(rgb_name)[:, :, :3][:, :, ::-1].copy() / 255.0
+        H, W = rgb.shape[0:2]
+        rgb_image, mask_image = rgb[:, :H], rgb[:, -H:, :1]
+        rgb_image = cv2.resize(rgb_image, IMAGE_SIZE)
+        rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device)
+
+        mask_image = cv2.resize(mask_image, IMAGE_SIZE).reshape(IMAGE_SIZE[0], IMAGE_SIZE[1], 1)
+        mask_image = torch.FloatTensor(mask_image).unsqueeze(0).to(guidance.device)
+
+        temp = torch.zeros(1).to(guidance.device)
+        azimuth = torch.tensor([rgb_idx/n_rgbs * AZIMUTH_FACTOR]).to(guidance.device)
+        camera_distance = torch.tensor([cfg.data.default_camera_distance]).to(guidance.device)
+
+        if cfg.data.view_dependent_noise:
+            guidance.min_step_percent = 0. + (rgb_idx/n_rgbs) * (cfg.system.guidance.min_step_percent)
+            guidance.max_step_percent = 0. + (rgb_idx/n_rgbs) * (cfg.system.guidance.max_step_percent)
+
+        denoised_image = process_guidance(cfg, guidance, prompt_utils, rgb_image, azimuth, temp, camera_distance, mask_image)
+        
+        save_image(denoised_image.permute(0,3,1,2), f"{cfg.data.save_path}/img_{azimuth[0]}.png", normalize=True, value_range=(0, 1))
+
+        copy_file(rgb_name.replace("png", "npy"), f"{cfg.data.save_path}/img_{azimuth[0]}.npy")
+
+        if rgb_idx == 0:
+            copy_file(rgb_name.replace("png", "npy"), f"{cfg.data.save_path}/ref_{azimuth[0]}.npy")       
+
+
+def process_guidance(cfg, guidance, prompt_utils, rgb_image, azimuth, temp, camera_distance, mask_image):
+    if cfg.data.azimuth_range[0] < azimuth < cfg.data.azimuth_range[1]:
+        return guidance.sample_img2img(
+            rgb_image, prompt_utils, temp, 
+            azimuth, camera_distance, seed=0, mask=mask_image
+        )["edit_image"]
+    else:
+        return rgb_image
+
+
+def generate_mv_dataset(cfg):
+
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(cfg.system.prompt_processor)
+    prompt_utils = prompt_processor()
+
+    guidance.update_step(epoch=0, global_step=0)
+    rgb_samples = prepare_images(cfg)
+    print(rgb_samples)
+    process_images(rgb_samples, cfg, guidance, prompt_utils)
+

threestudio/scripts/img_to_mv.py

@@ -0,0 +1,84 @@
+import os
+import argparse
+from PIL import Image
+import torch
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, StableDiffusionUpscalePipeline
+
+
+def load_model(superres):
+    mv_model = DiffusionPipeline.from_pretrained(
+        "sudo-ai/zero123plus-v1.1", custom_pipeline="sudo-ai/zero123plus-pipeline",
+        torch_dtype=torch.float16, cache_dir="load/checkpoints/huggingface/hub", local_files_only=True,
+    )
+    mv_model.scheduler = EulerAncestralDiscreteScheduler.from_config(
+        mv_model.scheduler.config, timestep_spacing='trailing', cache_dir="load/checkpoints/huggingface/hub", local_files_only=True,
+    )
+
+    if superres:
+        superres_model = StableDiffusionUpscalePipeline.from_pretrained(
+            "stabilityai/stable-diffusion-x4-upscaler", revision="fp16",
+            torch_dtype=torch.float16, cache_dir="load/checkpoints/huggingface/hub", local_files_only=True,
+        )
+    else:
+        superres_model = None
+
+    return mv_model, superres_model
+
+    
+def superres_4x(image, model, prompt):
+    low_res_img = image.resize((256, 256))
+    model.to('cuda:1')
+    result = model(prompt=prompt, image=low_res_img).images[0]
+    return result
+
+
+def img_to_mv(image_path, model):
+    cond = Image.open(image_path)
+    model.to('cuda:1')
+    result = model(cond, num_inference_steps=75).images[0]
+    return result
+
+
+def crop_save_image_to_2x3_grid(image, args, model):
+    save_path = args.save_path
+    width, height = image.size
+    grid_width = width//2
+    grid_height = height//3
+
+    images = []
+    for i in range(3):
+        for j in range(2):
+            left = j * grid_width
+            upper = i * grid_height
+            right = (j+1) * grid_width
+            lower = (i+1) * grid_height
+            
+            cropped_image = image.crop((left, upper, right, lower))
+            if args.superres:
+                cropped_image = superres_4x(cropped_image, model, args.prompt)
+            images.append(cropped_image)
+
+    for idx, img in enumerate(images):
+        img.save(os.path.join(save_path, f'cropped_{idx}.jpg'))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image_path', type=str, help="path to image (png, jpeg, etc.)")
+    parser.add_argument('--save_path', type=str, help="path to save output images")
+    parser.add_argument('--prompt', type=str, help="prompt to use for superres")
+    parser.add_argument('--superres', action='store_true', help="whether to use superres")
+    args = parser.parse_args()
+
+    print(args.superres)
+
+    os.makedirs(args.save_path, exist_ok=True)
+    os.system(f"cp '{args.image_path}' '{args.save_path}'")
+
+    mv_model, superres_model = load_model(args.superres)
+    images = img_to_mv(args.image_path, mv_model)
+    crop_save_image_to_2x3_grid(images, args, superres_model)
+
+
+# Example usage:
+# python threestudio/scripts/img_to_mv.py --image_path 'mushroom.png' --save_path '.cache/temp' --prompt 'a photo of mushroom' --superres

threestudio/scripts/make_training_vid.py

@@ -0,0 +1,77 @@
+# make_training_vid("outputs/zero123/64_teddy_rgba.png@20230627-195615", frames_per_vid=30, fps=20, max_iters=200)
+import argparse
+import glob
+import os
+
+import imageio
+import numpy as np
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+
+
+def draw_text_in_image(img, texts):
+    img = Image.fromarray(img)
+    draw = ImageDraw.Draw(img)
+    black, white = (0, 0, 0), (255, 255, 255)
+    for i, text in enumerate(texts):
+        draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
+        draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
+        draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
+        draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
+        draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
+    return np.asarray(img)
+
+
+def make_training_vid(exp, frames_per_vid=1, fps=3, max_iters=None, max_vids=None):
+    # exp = "/admin/home-vikram/git/threestudio/outputs/zero123/64_teddy_rgba.png@20230627-195615"
+    files = glob.glob(os.path.join(exp, "save", "*.mp4"))
+    if os.path.join(exp, "save", "training_vid.mp4") in files:
+        files.remove(os.path.join(exp, "save", "training_vid.mp4"))
+    its = [int(os.path.basename(file).split("-")[0].split("it")[-1]) for file in files]
+    it_sort = np.argsort(its)
+    files = list(np.array(files)[it_sort])
+    its = list(np.array(its)[it_sort])
+    max_vids = max_iters // its[0] if max_iters is not None else max_vids
+    files, its = files[:max_vids], its[:max_vids]
+    frames, i = [], 0
+    for it, file in tqdm(zip(its, files), total=len(files)):
+        vid = imageio.mimread(file)
+        for _ in range(frames_per_vid):
+            frame = vid[i % len(vid)]
+            frame = draw_text_in_image(frame, [str(it)])
+            frames.append(frame)
+            i += 1
+    # Save
+    imageio.mimwrite(os.path.join(exp, "save", "training_vid.mp4"), frames, fps=fps)
+
+
+def join(file1, file2, name):
+    # file1 = "/admin/home-vikram/git/threestudio/outputs/zero123/OLD_64_dragon2_rgba.png@20230629-023028/save/it200-val.mp4"
+    # file2 = "/admin/home-vikram/git/threestudio/outputs/zero123/64_dragon2_rgba.png@20230628-152734/save/it200-val.mp4"
+    vid1 = imageio.mimread(file1)
+    vid2 = imageio.mimread(file2)
+    frames = []
+    for f1, f2 in zip(vid1, vid2):
+        frames.append(
+            np.concatenate([f1[:, : f1.shape[0]], f2[:, : f2.shape[0]]], axis=1)
+        )
+    imageio.mimwrite(name, frames)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--exp", help="directory of experiment")
+    parser.add_argument(
+        "--frames_per_vid", type=int, default=1, help="# of frames from each val vid"
+    )
+    parser.add_argument("--fps", type=int, help="max # of iters to save")
+    parser.add_argument("--max_iters", type=int, help="max # of iters to save")
+    parser.add_argument(
+        "--max_vids",
+        type=int,
+        help="max # of val videos to save. Will be overridden by max_iters",
+    )
+    args = parser.parse_args()
+    make_training_vid(
+        args.exp, args.frames_per_vid, args.fps, args.max_iters, args.max_vids
+    )

threestudio/scripts/metric_utils.py

@@ -0,0 +1,459 @@
+# * evaluate use laion/CLIP-ViT-H-14-laion2B-s32B-b79K
+# best open source clip so far: laion/CLIP-ViT-bigG-14-laion2B-39B-b160k
+# code adapted from NeuralLift-360
+
+import torch
+import torch.nn as nn
+import os
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+import matplotlib.pyplot as plt
+# import clip
+from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTokenizer, CLIPProcessor
+from torchvision import transforms
+import numpy as np
+import torch.nn.functional as F
+from tqdm import tqdm
+import cv2
+from PIL import Image
+# import torchvision.transforms as transforms
+import glob
+from skimage.metrics import peak_signal_noise_ratio as compare_psnr
+import lpips
+from os.path import join as osp
+import argparse
+import pandas as pd
+
+class CLIP(nn.Module):
+
+    def __init__(self,
+                 device,
+                 clip_name='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k',
+                 size=224):  #'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'):
+        super().__init__()
+        self.size = size
+        self.device = f"cuda:{device}"
+
+        clip_name = clip_name
+
+        self.feature_extractor = CLIPFeatureExtractor.from_pretrained(
+            clip_name)
+        self.clip_model = CLIPModel.from_pretrained(clip_name).to(self.device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            'openai/clip-vit-base-patch32')
+
+        self.normalize = transforms.Normalize(
+            mean=self.feature_extractor.image_mean,
+            std=self.feature_extractor.image_std)
+
+        self.resize = transforms.Resize(224)
+        self.to_tensor = transforms.ToTensor()
+
+        # image augmentation
+        self.aug = T.Compose([
+            T.Resize((224, 224)),
+            T.Normalize((0.48145466, 0.4578275, 0.40821073),
+                        (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+    # * recommend to use this function for evaluation
+    @torch.no_grad()
+    def score_gt(self, ref_img_path, novel_views):
+        # assert len(novel_views) == 100
+        clip_scores = []
+        for novel in novel_views:
+            clip_scores.append(self.score_from_path(ref_img_path, [novel]))
+        return np.mean(clip_scores)
+
+    # * recommend to use this function for evaluation
+    # def score_gt(self, ref_paths, novel_paths):
+    #     clip_scores = []
+    #     for img1_path, img2_path in zip(ref_paths, novel_paths):
+    #         clip_scores.append(self.score_from_path(img1_path, img2_path))
+
+    #     return np.mean(clip_scores)
+
+    def similarity(self, image1_features: torch.Tensor,
+                   image2_features: torch.Tensor) -> float:
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            y = image1_features.T.view(image1_features.T.shape[1],
+                                       image1_features.T.shape[0])
+            similarity = torch.matmul(y, image2_features.T)
+            # print(similarity)
+            return similarity[0][0].item()
+
+    def get_img_embeds(self, img):
+        if img.shape[0] == 4:
+            img = img[:3, :, :]
+
+        img = self.aug(img).to(self.device)
+        img = img.unsqueeze(0)  # b,c,h,w
+
+        # plt.imshow(img.cpu().squeeze(0).permute(1, 2, 0).numpy())
+        # plt.show()
+        # print(img)
+
+        image_z = self.clip_model.get_image_features(img)
+        image_z = image_z / image_z.norm(dim=-1,
+                                         keepdim=True)  # normalize features
+        return image_z
+
+    def score_from_feature(self, img1, img2):
+        img1_feature, img2_feature = self.get_img_embeds(
+            img1), self.get_img_embeds(img2)
+        # for debug
+        return self.similarity(img1_feature, img2_feature)
+
+    def read_img_list(self, img_list):
+        size = self.size
+        images = []
+        # white_background = np.ones((size, size, 3), dtype=np.uint8) * 255
+
+        for img_path in img_list:
+            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+            # print(img_path)
+            if img.shape[2] == 4:  # Handle BGRA images
+                alpha = img[:, :, 3]  # Extract alpha channel
+                img = cv2.cvtColor(img,cv2.COLOR_BGRA2RGB)  # Convert BGRA to BGR
+                img[np.where(alpha == 0)] = [
+                    255, 255, 255
+                ]  # Set transparent pixels to white
+            else:  # Handle other image formats like JPG and PNG
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+            img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+
+            # plt.imshow(img)
+            # plt.show()
+
+            images.append(img)
+
+        images = np.stack(images, axis=0)
+        # images[np.where(images == 0)] = 255  # Set black pixels to white
+        # images = np.where(images == 0, white_background, images)  # Set transparent pixels to white
+        # images = images.astype(np.float32)
+
+        return images
+
+    def score_from_path(self, img1_path, img2_path):
+        img1, img2 = self.read_img_list(img1_path), self.read_img_list(img2_path)
+        img1 = np.squeeze(img1)
+        img2 = np.squeeze(img2)
+        # plt.imshow(img1)
+        # plt.show()
+        # plt.imshow(img2)
+        # plt.show()
+
+        img1, img2 = self.to_tensor(img1), self.to_tensor(img2)
+        # print("img1 to tensor ",img1)
+        return self.score_from_feature(img1, img2)
+
+
+def numpy_to_torch(images):
+    images = images * 2.0 - 1.0
+    images = torch.from_numpy(images.transpose((0, 3, 1, 2))).float()
+    return images.cuda()
+
+
+class LPIPSMeter:
+
+    def __init__(self,
+                 net='alex',
+                 device=None,
+                 size=224):  # or we can use 'alex', 'vgg' as network
+        self.size = size
+        self.net = net
+        self.results = []
+        self.device = device if device is not None else torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
+        self.fn = lpips.LPIPS(net=net).eval().to(self.device)
+
+    def measure(self):
+        return np.mean(self.results)
+
+    def report(self):
+        return f'LPIPS ({self.net}) = {self.measure():.6f}'
+
+    def read_img_list(self, img_list):
+        size = self.size
+        images = []
+        white_background = np.ones((size, size, 3), dtype=np.uint8) * 255
+
+        for img_path in img_list:
+            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+
+            if img.shape[2] == 4:  # Handle BGRA images
+                alpha = img[:, :, 3]  # Extract alpha channel
+                img = cv2.cvtColor(img,
+                                   cv2.COLOR_BGRA2BGR)  # Convert BGRA to BGR
+
+                img = cv2.cvtColor(img,
+                                   cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
+                img[np.where(alpha == 0)] = [
+                    255, 255, 255
+                ]  # Set transparent pixels to white
+            else:  # Handle other image formats like JPG and PNG
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+            img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+            images.append(img)
+
+        images = np.stack(images, axis=0)
+        # images[np.where(images == 0)] = 255  # Set black pixels to white
+        # images = np.where(images == 0, white_background, images)  # Set transparent pixels to white
+        images = images.astype(np.float32) / 255.0
+
+        return images
+
+    # * recommend to use this function for evaluation
+    @torch.no_grad()
+    def score_gt(self, ref_paths, novel_paths):
+        self.results = []
+        for path0, path1 in zip(ref_paths, novel_paths):
+            # Load images
+            # img0 = lpips.im2tensor(lpips.load_image(path0)).cuda() # RGB image from [-1,1]
+            # img1 = lpips.im2tensor(lpips.load_image(path1)).cuda()
+            img0, img1 = self.read_img_list([path0]), self.read_img_list(
+                [path1])
+            img0, img1 = numpy_to_torch(img0), numpy_to_torch(img1)
+            # print(img0.shape,img1.shape)
+            img0 = F.interpolate(img0,
+                                    size=(self.size, self.size),
+                                    mode='area')
+            img1 = F.interpolate(img1,
+                                    size=(self.size, self.size),
+                                    mode='area')
+
+            # for debug vis
+            # plt.imshow(img0.cpu().squeeze(0).permute(1, 2, 0).numpy())
+            # plt.show()
+            # plt.imshow(img1.cpu().squeeze(0).permute(1, 2, 0).numpy())
+            # plt.show()
+            # equivalent to cv2.resize(rgba, (w, h), interpolation=cv2.INTER_AREA
+
+            # print(img0.shape,img1.shape)
+
+            self.results.append(self.fn.forward(img0, img1).cpu().numpy())
+
+        return self.measure()
+
+
+class PSNRMeter:
+
+    def __init__(self, size=800):
+        self.results = []
+        self.size = size
+
+    def read_img_list(self, img_list):
+        size = self.size
+        images = []
+        white_background = np.ones((size, size, 3), dtype=np.uint8) * 255
+        for img_path in img_list:
+            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+
+            if img.shape[2] == 4:  # Handle BGRA images
+                alpha = img[:, :, 3]  # Extract alpha channel
+                img = cv2.cvtColor(img,
+                                   cv2.COLOR_BGRA2BGR)  # Convert BGRA to BGR
+
+                img = cv2.cvtColor(img,
+                                   cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
+                img[np.where(alpha == 0)] = [
+                    255, 255, 255
+                ]  # Set transparent pixels to white
+            else:  # Handle other image formats like JPG and PNG
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+            img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+            images.append(img)
+
+        images = np.stack(images, axis=0)
+        # images[np.where(images == 0)] = 255  # Set black pixels to white
+        # images = np.where(images == 0, white_background, images)  # Set transparent pixels to white
+        images = images.astype(np.float32) / 255.0
+        # print(images.shape)
+        return images
+
+    def update(self, preds, truths):
+        # print(preds.shape)
+
+        psnr_values = []
+        # For each pair of images in the batches
+        for img1, img2 in zip(preds, truths):
+            # Compute the PSNR and add it to the list
+            # print(img1.shape,img2.shape)
+
+            # for debug
+            # plt.imshow(img1)
+            # plt.show()
+            # plt.imshow(img2)
+            # plt.show()
+
+            psnr = compare_psnr(
+                img1, img2,
+                data_range=1.0)  # assuming your images are scaled to [0,1]
+            # print(f"temp psnr {psnr}")
+            psnr_values.append(psnr)
+
+        # Convert the list of PSNR values to a numpy array
+        self.results = psnr_values
+
+    def measure(self):
+        return np.mean(self.results)
+
+    def report(self):
+        return f'PSNR = {self.measure():.6f}'
+
+    # * recommend to use this function for evaluation
+    def score_gt(self, ref_paths, novel_paths):
+        self.results = []
+        # [B, N, 3] or [B, H, W, 3], range[0, 1]
+        preds = self.read_img_list(ref_paths)
+        truths = self.read_img_list(novel_paths)
+        self.update(preds, truths)
+        return self.measure()
+
+all_inputs = 'data'
+nerf_dataset = os.listdir(osp(all_inputs, 'nerf4'))
+realfusion_dataset = os.listdir(osp(all_inputs, 'realfusion15'))
+meta_examples = {
+   'nerf4': nerf_dataset, 
+   'realfusion15': realfusion_dataset, 
+}
+all_datasets = meta_examples.keys()
+
+# organization 1
+def deprecated_score_from_method_for_dataset(my_scorer,
+                                  method,
+                                  dataset,
+                                  input,
+                                  output,
+                                  score_type='clip', 
+                                  ):  # psnr, lpips
+    # print("\n\n\n")
+    # print(f"______{method}___{dataset}___{score_type}_________")
+    scores = {}
+    final_res = 0
+    examples = meta_examples[dataset]
+    for i in range(len(examples)):
+        
+        # compare entire folder for clip
+        if score_type == 'clip':
+            novel_view = osp(pred_path, examples[i], 'colors')
+        # compare first image for other metrics
+        else:
+            if method == '3d_fuse': method = '3d_fuse_0'
+            novel_view = list(
+                glob.glob(
+                    osp(pred_path, examples[i], 'colors',
+                        'step_0000*')))[0]
+
+        score_i = my_scorer.score_gt(
+            [], [novel_view])
+        scores[examples[i]] = score_i
+        final_res += score_i
+    # print(scores, " Avg : ", final_res / len(examples))
+    # print("``````````````````````")
+    return scores
+
+# results organization 2
+def score_from_method_for_dataset(my_scorer,
+                                  input_path,
+                                  pred_path,
+                                  score_type='clip', 
+                                  rgb_name='lambertian', 
+                                  result_folder='results/images', 
+                                  first_str='*0000*'
+                                  ):  # psnr, lpips
+    scores = {}
+    final_res = 0
+    examples = os.listdir(input_path)
+    for i in range(len(examples)):
+        # ref path
+        ref_path = osp(input_path, examples[i], 'rgba.png') 
+        # compare entire folder for clip
+        if score_type == 'clip':
+            novel_view = glob.glob(osp(pred_path,'*'+examples[i]+'*', result_folder, f'*{rgb_name}*'))
+            print(f'[INOF] {score_type} loss for example {examples[i]} between 1 GT and {len(novel_view)} predictions')
+        # compare first image for other metrics
+        else:
+            novel_view = glob.glob(osp(pred_path, '*'+examples[i]+'*/', result_folder, f'{first_str}{rgb_name}*'))
+            print(f'[INOF] {score_type} loss for example {examples[i]} between {ref_path} and {novel_view}')
+        # breakpoint()
+        score_i = my_scorer.score_gt([ref_path], novel_view)
+        scores[examples[i]] = score_i
+        final_res += score_i
+    avg_score = final_res / len(examples)
+    scores['average'] = avg_score
+    return scores
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Script to accept three string arguments")
+    parser.add_argument("--input_path",
+                        default=all_inputs,
+                        help="Specify the input path")
+    parser.add_argument("--pred_pattern",
+                        default="out/magic123*", 
+                        help="Specify the pattern of predition paths")
+    parser.add_argument("--results_folder",
+                        default="results/images", 
+                        help="where are the results under each pred_path")
+    parser.add_argument("--rgb_name",
+                        default="lambertian", 
+                        help="the postfix of the image")
+    parser.add_argument("--first_str",
+                        default="*0000*", 
+                        help="the str to indicate the first view")
+    parser.add_argument("--datasets",
+                        default=all_datasets,
+                        nargs='*',
+                        help="Specify the output path")
+    parser.add_argument("--device",
+                        type=int,
+                        default=0,
+                        help="Specify the GPU device to be used")
+    parser.add_argument("--save_dir", type=str, default='all_metrics/results')
+    args = parser.parse_args()
+
+    clip_scorer = CLIP(args.device)
+    lpips_scorer = LPIPSMeter()
+    psnr_scorer = PSNRMeter()
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    for dataset in args.datasets:
+        input_path = osp(args.input_path, dataset)
+        
+        # assume the pred_path is organized as: pred_path/methods/dataset
+        pred_pattern = osp(args.pred_pattern, dataset)
+        pred_paths = glob.glob(pred_pattern)
+        print(f"[INFO] Following the pattern {pred_pattern}, find {len(pred_paths)} pred_paths: \n", pred_paths)
+        if len(pred_paths) == 0:
+            raise IOError
+        for pred_path in pred_paths:
+            if not os.path.exists(pred_path):
+                print(f'[WARN] prediction does not exit for {pred_path}')
+            else:
+                print(f'[INFO] evaluate {pred_path}')
+            results_dict = {}
+            results_dict['clip'] = score_from_method_for_dataset(
+                clip_scorer, input_path, pred_path, 'clip', 
+                result_folder=args.results_folder, rgb_name=args.rgb_name, first_str=args.first_str)
+            
+            results_dict['psnr'] = score_from_method_for_dataset(
+                psnr_scorer, input_path, pred_path,  'psnr', 
+                result_folder=args.results_folder, rgb_name=args.rgb_name, first_str=args.first_str)
+            
+            results_dict['lpips'] = score_from_method_for_dataset(
+                lpips_scorer, input_path, pred_path,  'lpips', 
+                result_folder=args.results_folder, rgb_name=args.rgb_name, first_str=args.first_str)
+            
+            df = pd.DataFrame(results_dict)
+            method = pred_path.split('/')[-2]
+            print(osp(pred_path, args.results_folder))
+            results_str = '_'.join(args.results_folder.split('/'))
+            print(method+'-'+results_str)
+            print(df)
+            df.to_csv(f"{args.save_dir}/{method}-{results_str}-{dataset}.csv")

threestudio/scripts/run_gaussian.sh

@@ -0,0 +1,36 @@
+import subprocess
+
+prompt_list = [
+    "a delicious hamburger",
+    "A DSLR photo of a roast turkey on a platter",
+    "A high quality photo of a dragon",
+    "A DSLR photo of a bald eagle",
+    "A bunch of blue rose, highly detailed",
+    "A 3D model of an adorable cottage with a thatched roof",
+    "A high quality photo of a furry corgi",
+    "A DSLR photo of a panda",
+    "a DSLR photo of a cat lying on its side batting at a ball of yarn",
+    "a beautiful dress made out of fruit, on a mannequin. Studio lighting, high quality, high resolution",
+    "a DSLR photo of a corgi wearing a beret and holding a baguette, standing up on two hind legs",
+    "a zoomed out DSLR photo of a stack of pancakes",
+    "a zoomed out DSLR photo of a baby bunny sitting on top of a stack of pancakes",
+]
+negative_prompt = "oversaturated color, ugly, tiling, low quality, noise, ugly pattern"
+
+gpu_id = 0
+max_steps = 10
+val_check = 1
+out_name = "gsgen_baseline"
+for prompt in prompt_list:
+    print(f"Running model on device {gpu_id}: ", prompt)
+    command = [
+        "python", "launch.py",
+        "--config", "configs/gaussian_splatting.yaml",
+        "--train",
+        f"system.prompt_processor.prompt={prompt}",
+        f"system.prompt_processor.negative_prompt={negative_prompt}",
+        f"name={out_name}",
+        "--gpu", f"{gpu_id}"
+    ]
+    subprocess.run(command)
+        

threestudio/scripts/run_zero123.py

@@ -0,0 +1,13 @@
+NAME="dragon2"
+
+# Phase 1 - 64x64
+python launch.py --config configs/zero123.yaml --train --gpu 7 data.image_path=./load/images/${NAME}_rgba.png use_timestamp=False name=${NAME} tag=Phase1 # system.freq.guidance_eval=0 system.loggers.wandb.enable=false system.loggers.wandb.project="zero123" system.loggers.wandb.name=${NAME}_Phase1
+
+# Phase 1.5 - 512 refine
+python launch.py --config configs/zero123-geometry.yaml --train --gpu 4 data.image_path=./load/images/${NAME}_rgba.png system.geometry_convert_from=./outputs/${NAME}/Phase1/ckpts/last.ckpt use_timestamp=False name=${NAME} tag=Phase1p5 # system.freq.guidance_eval=0 system.loggers.wandb.enable=false system.loggers.wandb.project="zero123" system.loggers.wandb.name=${NAME}_Phase1p5
+
+# Phase 2 - dreamfusion
+python launch.py --config configs/experimental/imagecondition_zero123nerf.yaml --train --gpu 5 data.image_path=./load/images/${NAME}_rgba.png system.prompt_processor.prompt="A 3D model of a friendly dragon" system.weights="/admin/home-vikram/git/threestudio/outputs/${NAME}/Phase1/ckpts/last.ckpt" name=${NAME} tag=Phase2 # system.freq.guidance_eval=0 system.loggers.wandb.enable=false system.loggers.wandb.project="zero123" system.loggers.wandb.name=${NAME}_Phase2
+
+# Phase 2 - SDF + dreamfusion
+python launch.py --config configs/experimental/imagecondition_zero123nerf_refine.yaml --train --gpu 5 data.image_path=./load/images/${NAME}_rgba.png system.prompt_processor.prompt="A 3D model of a friendly dragon" system.geometry_convert_from="/admin/home-vikram/git/threestudio/outputs/${NAME}/Phase1/ckpts/last.ckpt" name=${NAME} tag=Phase2_refine # system.freq.guidance_eval=0 system.loggers.wandb.enable=false system.loggers.wandb.project="zero123" system.loggers.wandb.name=${NAME}_Phase2_refine

threestudio/scripts/run_zero123_comparison.sh

@@ -0,0 +1,23 @@
+# with standard zero123
+threestudio/scripts/run_zero123_phase.sh 6 anya_front 105000 0
+
+# with zero123XL (not released yet!)
+threestudio/scripts/run_zero123_phase.sh 1 anya_front XL_20230604 0
+threestudio/scripts/run_zero123_phase.sh 2 baby_phoenix_on_ice XL_20230604 20
+threestudio/scripts/run_zero123_phase.sh 3 beach_house_1 XL_20230604 50
+threestudio/scripts/run_zero123_phase.sh 4 bollywood_actress XL_20230604 0
+threestudio/scripts/run_zero123_phase.sh 5 beach_house_2 XL_20230604 30
+threestudio/scripts/run_zero123_phase.sh 6 hamburger XL_20230604 10
+threestudio/scripts/run_zero123_phase.sh 7 cactus XL_20230604 8
+threestudio/scripts/run_zero123_phase.sh 0 catstatue XL_20230604 50
+threestudio/scripts/run_zero123_phase.sh 1 church_ruins XL_20230604 0
+threestudio/scripts/run_zero123_phase.sh 2 firekeeper XL_20230604 10
+threestudio/scripts/run_zero123_phase.sh 3 futuristic_car XL_20230604 20
+threestudio/scripts/run_zero123_phase.sh 4 mona_lisa XL_20230604 10
+threestudio/scripts/run_zero123_phase.sh 5 teddy XL_20230604 20
+
+# set guidance_eval to 0, to greatly speed up training
+threestudio/scripts/run_zero123_phase.sh 7 anya_front XL_20230604 0 system.freq.guidance_eval=0
+
+# disable wandb for faster training (or if you don't want to use it)
+threestudio/scripts/run_zero123_phase.sh 7 anya_front XL_20230604 0 system.loggers.wandb.enable=false system.freq.guidance_eval=0

threestudio/scripts/run_zero123_demo.sh

@@ -0,0 +1,25 @@
+NAME="dragon2"
+
+# Phase 1 - 64x64
+python launch.py --config configs/zero123_64.yaml --train --gpu 7 system.loggers.wandb.enable=false system.loggers.wandb.project="voletiv-anya-new" system.loggers.wandb.name=${NAME} data.image_path=./load/images/${NAME}_rgba.png system.freq.guidance_eval=0 system.guidance.pretrained_model_name_or_path="./load/zero123/XL_20230604.ckpt" use_timestamp=False name=${NAME} tag="Phase1_64"
+
+# python threestudio/scripts/make_training_vid.py --exp /admin/home-vikram/git/threestudio/outputs/zero123/64_dragon2_rgba.png@20230628-152734 --frames_per_vid 30 --fps 20 --max_iters 200
+
+# # Phase 1.5 - 512
+# python launch.py --config configs/zero123_512.yaml --train --gpu 5 system.loggers.wandb.enable=true system.loggers.wandb.project="voletiv-zero123XL-demo" system.loggers.wandb.name="robot_512_drel_n_XL_SAMEgeom" data.image_path=./load/images/robot_rgba.png system.freq.guidance_eval=0 system.guidance.pretrained_model_name_or_path="./load/zero123/XL_20230604.ckpt" tag='${data.random_camera.height}_${rmspace:${basename:${data.image_path}},_}_XL_SAMEgeom' system.weights="/admin/home-vikram/git/threestudio/outputs/zero123/[64, 128]_robot_rgba.png_OLD@20230630-052314/ckpts/last.ckpt"
+
+# Phase 1.5 - 512 refine
+python launch.py --config configs/zero123-geometry.yaml --train --gpu 4 system.loggers.wandb.enable=false system.loggers.wandb.project="voletiv-zero123XL-demo" system.loggers.wandb.name="robot_512_drel_n_XL_SAMEg" system.freq.guidance_eval=0 data.image_path=./load/images/${NAME}_rgba.png system.geometry_convert_from=./outputs/${NAME}/Phase1_64/ckpts/last.ckpt use_timestamp=False name=${NAME} tag="Phase2_512geom"
+
+# Phase 2 - dreamfusion
+python launch.py --config configs/experimental/imagecondition_zero123nerf.yaml --train --gpu 5 system.loggers.wandb.enable=false system.loggers.wandb.project="voletiv-zero123XL-demo" system.loggers.wandb.name="robot_512_drel_n_XL_SAMEw" tag='${data.random_camera.height}_${rmspace:${basename:${data.image_path}},_}_XL_Phase2' system.freq.guidance_eval=0 data.image_path=./load/images/robot_rgba.png system.prompt_processor.prompt="A DSLR 3D photo of a cute anime schoolgirl stands proudly with her arms in the air, pink hair ( unreal engine 5 trending on Artstation Ghibli 4k )" system.weights="/admin/home-vikram/git/threestudio/outputs/zero123/[64, 128]_robot_rgba.png_OLD@20230630-052314/ckpts/last.ckpt"
+
+python launch.py --config configs/experimental/imagecondition_zero123nerf_refine.yaml --train --gpu 5 system.loggers.wandb.enable=false system.loggers.wandb.project="voletiv-zero123XL-demo" system.loggers.wandb.name="robot_512_drel_n_XL_SAMEw" tag='${data.random_camera.height}_${rmspace:${basename:${data.image_path}},_}_XL_Phase2_refine' system.freq.guidance_eval=0 data.image_path=./load/images/robot_rgba.png system.prompt_processor.prompt="A 3D model of a friendly dragon" system.geometry_convert_from="/admin/home-vikram/git/threestudio/outputs/zero123/[64, 128, 256]_dragon2_rgba.png_XL_REPEAT@20230705-023531/ckpts/last.ckpt"
+
+# A DSLR 3D photo of a cute anime schoolgirl stands proudly with her arms in the air, pink hair ( unreal engine 5 trending on Artstation Ghibli 4k )"
+# "/admin/home-vikram/git/threestudio/outputs/zero123/[64, 128]_robot_rgba.png_OLD@20230630-052314/ckpts/last.ckpt"
+
+# Adds zero123_512-refine.yaml
+# Adds resolution_milestones to image.py
+# guidance_eval gets max batch_size 4
+# Introduces random_bg in solid_color_bg

threestudio/scripts/run_zero123_phase.sh

@@ -0,0 +1,14 @@
+
+GPU_ID=$1         # e.g. 0
+IMAGE_PREFIX=$2   # e.g. "anya_front"
+ZERO123_PREFIX=$3 # e.g. "XL_20230604"
+ELEVATION=$4      # e.g. 0
+REST=${@:5:99}    # e.g. "system.guidance.min_step_percent=0.1 system.guidance.max_step_percent=0.9"
+
+# change this config if you don't use wandb or want to speed up training
+python launch.py --config configs/zero123.yaml --train --gpu $GPU_ID system.loggers.wandb.enable=true system.loggers.wandb.project="claforte-noise_atten" \
+    system.loggers.wandb.name="${IMAGE_PREFIX}_zero123_${ZERO123_PREFIX}...fov20_${REST}" \
+    data.image_path=./load/images/${IMAGE_PREFIX}_rgba.png system.freq.guidance_eval=37 \
+    system.guidance.pretrained_model_name_or_path="./load/zero123/${ZERO123_PREFIX}.ckpt" \
+    system.guidance.cond_elevation_deg=$ELEVATION \
+    ${REST}

threestudio/scripts/run_zero123_phase2.sh

@@ -0,0 +1,5 @@
+# Reconstruct Anya using latest Zero123XL, in <2000 steps.
+python launch.py --config configs/zero123.yaml --train --gpu 0 system.loggers.wandb.enable=true system.loggers.wandb.project="voletiv-anya-new" system.loggers.wandb.name="claforte_params" data.image_path=./load/images/anya_front_rgba.png system.freq.ref_or_zero123="accumulate" system.freq.guidance_eval=13 system.guidance.pretrained_model_name_or_path="./load/zero123/XL_20230604.ckpt"
+
+# PHASE 2
+python launch.py --config configs/experimental/imagecondition_zero123nerf.yaml --train --gpu 0 system.prompt_processor.prompt="A DSLR 3D photo of a cute anime schoolgirl stands proudly with her arms in the air, pink hair ( unreal engine 5 trending on Artstation Ghibli 4k )" system.weights=outputs/zero123/128_anya_front_rgba.png@20230623-145711/ckpts/last.ckpt system.freq.guidance_eval=13 system.loggers.wandb.enable=true system.loggers.wandb.project="voletiv-anya-new" data.image_path=./load/images/anya_front_rgba.png system.loggers.wandb.name="anya" data.random_camera.progressive_until=500

threestudio/scripts/test_dreambooth.py

@@ -0,0 +1,54 @@
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+import torch
+
+# model_id = "load/checkpoints/sd_21_base_mushroom_vd_prompt"
+# model_id = "load/checkpoints/sd_base_mushroom"
+model_id = ".cache/checkpoints/sd_21_base_rabbit"
+# scheduler = DDIMScheduler()
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
+guidance_scale = 7.5
+
+prompt = "a sks rabbit, front view"
+image = pipe(prompt, num_inference_steps=50, guidance_scale=guidance_scale).images[0]
+
+image.save("debug.png")
+
+
+# import os
+# import cv2
+# import glob
+# import torch
+# import argparse
+# import numpy as np
+# from tqdm import tqdm
+# import pytorch_lightning as pl
+# from torchvision.utils import save_image
+
+# import threestudio
+# from threestudio.utils.config import load_config
+
+
+# if __name__ == "__main__":
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument("--config", required=True, help="path to config file")
+#     parser.add_argument("--view_dependent_noise", action="store_true", help="use view depdendent noise strength")
+
+#     args, extras = parser.parse_known_args()
+
+#     cfg = load_config(args.config, cli_args=extras, n_gpus=1)
+#     guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+#     prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(cfg.system.prompt_processor)
+#     prompt_utils = prompt_processor()
+
+#     guidance.update_step(epoch=0, global_step=0)
+#     elevation, azimuth = torch.zeros(1).cuda(), torch.zeros(1).cuda()
+#     camera_distances = torch.tensor([3.0]).cuda()
+#     c2w = torch.zeros(4,4).cuda()
+#     a = guidance.sample(prompt_utils, elevation, azimuth, camera_distances) # sample_lora
+#     from torchvision.utils import save_image
+#     save_image(a.permute(0,3,1,2), "debug.png", normalize=True, value_range=(0,1))
+
+
+
+# python threestudio/scripts/test_dreambooth.py --config configs/experimental/stablediffusion.yaml system.prompt_processor.prompt="a sks mushroom growing on a log" \
+#     system.guidance.pretrained_model_name_or_path_lora="load/checkpoints/sd_21_base_mushroom_camera_condition"

threestudio/scripts/test_dreambooth_lora.py

@@ -0,0 +1,25 @@
+import torch
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+
+
+# model_base = "stabilityai/stable-diffusion-2-1-base"
+
+# pipe = DiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16, cache_dir=CACHE_DIR, local_files_only=True)
+# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, cache_dir=CACHE_DIR, local_files_only=True)
+# lora_model_path = "load/checkpoints/sd_21_base_bear_dreambooth_lora"
+# pipe.unet.load_attn_procs(lora_model_path)
+
+# pipe.to("cuda")
+
+
+# image = pipe("A picture of a sks bear in the sky", num_inference_steps=50, guidance_scale=7.5).images[0]
+# image.save("bear_dreambooth_lora.png")
+
+
+pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", local_files_only=True, safety_checker=None)
+pipe.load_lora_weights("if_dreambooth_mushroom")
+pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")
+pipe.to("cuda:7")
+
+image = pipe("A photo of a sks mushroom, front view", num_inference_steps=50, guidance_scale=7.5).images[0]
+image.save("mushroom_dreambooth_lora.png")

threestudio/scripts/train_text_to_image_lora.py

@@ -0,0 +1,927 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Fine-tuning script for Stable Diffusion for text2image with support for LoRA."""
+
+import argparse
+import logging
+import math
+import os
+import random
+import shutil
+from pathlib import Path
+
+import datasets
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from datasets import load_dataset
+from huggingface_hub import create_repo, upload_folder
+from packaging import version
+from torchvision import transforms
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+
+import diffusers
+from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import compute_snr
+from diffusers.utils import check_min_version, is_wandb_available
+from diffusers.utils.import_utils import is_xformers_available
+
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.24.0.dev0")
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+def save_model_card(repo_id: str, images=None, base_model=str, dataset_name=str, repo_folder=None):
+    img_str = ""
+    for i, image in enumerate(images):
+        image.save(os.path.join(repo_folder, f"image_{i}.png"))
+        img_str += f"![img_{i}](./image_{i}.png)\n"
+
+    yaml = f"""
+---
+license: creativeml-openrail-m
+base_model: {base_model}
+tags:
+- stable-diffusion
+- stable-diffusion-diffusers
+- text-to-image
+- diffusers
+- lora
+inference: true
+---
+    """
+    model_card = f"""
+# LoRA text2image fine-tuning - {repo_id}
+These are LoRA adaption weights for {base_model}. The weights were fine-tuned on the {dataset_name} dataset. You can find some example images in the following. \n
+{img_str}
+"""
+    with open(os.path.join(repo_folder, "README.md"), "w") as f:
+        f.write(yaml + model_card)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help=(
+            "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
+            " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
+            " or to a folder containing files that 🤗 Datasets can understand."
+        ),
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The config of the Dataset, leave as None if there's only one config.",
+    )
+    parser.add_argument(
+        "--train_data_dir",
+        type=str,
+        default=None,
+        help=(
+            "A folder containing the training data. Folder contents must follow the structure described in"
+            " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
+            " must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
+        ),
+    )
+    parser.add_argument(
+        "--image_column", type=str, default="image", help="The column of the dataset containing an image."
+    )
+    parser.add_argument(
+        "--caption_column",
+        type=str,
+        default="text",
+        help="The column of the dataset containing a caption or a list of captions.",
+    )
+    parser.add_argument(
+        "--validation_prompt", type=str, default=None, help="A prompt that is sampled during training for inference."
+    )
+    parser.add_argument(
+        "--num_validation_images",
+        type=int,
+        default=4,
+        help="Number of images that should be generated during validation with `validation_prompt`.",
+    )
+    parser.add_argument(
+        "--validation_epochs",
+        type=int,
+        default=1,
+        help=(
+            "Run fine-tuning validation every X epochs. The validation process consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_samples",
+        type=int,
+        default=None,
+        help=(
+            "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        ),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="sd-model-finetuned-lora",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        type=str,
+        default=None,
+        help="The directory where the downloaded models and datasets will be stored.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop",
+        default=False,
+        action="store_true",
+        help=(
+            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+            " cropped. The images will be resized to the resolution first before cropping."
+        ),
+    )
+    parser.add_argument(
+        "--random_flip",
+        action="store_true",
+        help="whether to randomly flip images horizontally",
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--snr_gamma",
+        type=float,
+        default=None,
+        help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. "
+        "More details here: https://arxiv.org/abs/2303.09556.",
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--prediction_type",
+        type=str,
+        default=None,
+        help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
+    )
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
+            " training using `--resume_from_checkpoint`."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
+    )
+    parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.")
+    parser.add_argument(
+        "--rank",
+        type=int,
+        default=4,
+        help=("The dimension of the LoRA update matrices."),
+    )
+
+    args = parser.parse_args()
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_data_dir is None:
+        raise ValueError("Need either a dataset name or a training folder.")
+
+    return args
+
+
+DATASET_NAME_MAPPING = {
+    "lambdalabs/pokemon-blip-captions": ("image", "text"),
+}
+
+
+def main():
+    args = parse_args()
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+        import wandb
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+        if args.push_to_hub:
+            repo_id = create_repo(
+                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
+            ).repo_id
+    # Load scheduler, tokenizer and models.
+    noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
+    tokenizer = CLIPTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
+    )
+    text_encoder = CLIPTextModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+    )
+    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+    unet = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
+    )
+    # freeze parameters of models to save more memory
+    unet.requires_grad_(False)
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+
+    # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # Move unet, vae and text_encoder to device and cast to weight_dtype
+    unet.to(accelerator.device, dtype=weight_dtype)
+    vae.to(accelerator.device, dtype=weight_dtype)
+    text_encoder.to(accelerator.device, dtype=weight_dtype)
+
+    # now we will add new LoRA weights to the attention layers
+    # It's important to realize here how many attention weights will be added and of which sizes
+    # The sizes of the attention layers consist only of two different variables:
+    # 1) - the "hidden_size", which is increased according to `unet.config.block_out_channels`.
+    # 2) - the "cross attention size", which is set to `unet.config.cross_attention_dim`.
+
+    # Let's first see how many attention processors we will have to set.
+    # For Stable Diffusion, it should be equal to:
+    # - down blocks (2x attention layers) * (2x transformer layers) * (3x down blocks) = 12
+    # - mid blocks (2x attention layers) * (1x transformer layers) * (1x mid blocks) = 2
+    # - up blocks (2x attention layers) * (3x transformer layers) * (3x down blocks) = 18
+    # => 32 layers
+
+    # Set correct lora layers
+    lora_attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+
+        lora_attn_procs[name] = LoRAAttnProcessor(
+            hidden_size=hidden_size,
+            cross_attention_dim=cross_attention_dim,
+            rank=args.rank,
+        )
+
+    unet.set_attn_processor(lora_attn_procs)
+
+    if args.enable_xformers_memory_efficient_attention:
+        if is_xformers_available():
+            import xformers
+
+            xformers_version = version.parse(xformers.__version__)
+            if xformers_version == version.parse("0.0.16"):
+                logger.warn(
+                    "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+                )
+            unet.enable_xformers_memory_efficient_attention()
+        else:
+            raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+    lora_layers = AttnProcsLayers(unet.attn_processors)
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+        )
+
+    # Initialize the optimizer
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+    optimizer = optimizer_cls(
+        lora_layers.parameters(),
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    # Get the datasets: you can either provide your own training and evaluation files (see below)
+    # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub).
+
+    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        dataset = load_dataset(
+            args.dataset_name,
+            args.dataset_config_name,
+            cache_dir=args.cache_dir,
+            data_dir=args.train_data_dir,
+        )
+    else:
+        data_files = {}
+        if args.train_data_dir is not None:
+            data_files["train"] = os.path.join(args.train_data_dir, "**")
+        dataset = load_dataset(
+            "imagefolder",
+            data_files=data_files,
+            cache_dir=args.cache_dir,
+        )
+        # See more about loading custom images at
+        # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder
+
+    # Preprocessing the datasets.
+    # We need to tokenize inputs and targets.
+    column_names = dataset["train"].column_names
+
+    # 6. Get the column names for input/target.
+    dataset_columns = DATASET_NAME_MAPPING.get(args.dataset_name, None)
+    if args.image_column is None:
+        image_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
+    else:
+        image_column = args.image_column
+        if image_column not in column_names:
+            raise ValueError(
+                f"--image_column' value '{args.image_column}' needs to be one of: {', '.join(column_names)}"
+            )
+    if args.caption_column is None:
+        caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
+    else:
+        caption_column = args.caption_column
+        if caption_column not in column_names:
+            raise ValueError(
+                f"--caption_column' value '{args.caption_column}' needs to be one of: {', '.join(column_names)}"
+            )
+
+    # Preprocessing the datasets.
+    # We need to tokenize input captions and transform the images.
+    def tokenize_captions(examples, is_train=True):
+        captions = []
+        for caption in examples[caption_column]:
+            if isinstance(caption, str):
+                captions.append(caption)
+            elif isinstance(caption, (list, np.ndarray)):
+                # take a random caption if there are multiple
+                captions.append(random.choice(caption) if is_train else caption[0])
+            else:
+                raise ValueError(
+                    f"Caption column `{caption_column}` should contain either strings or lists of strings."
+                )
+        inputs = tokenizer(
+            captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
+        )
+        return inputs.input_ids
+
+    # Preprocessing the datasets.
+    train_transforms = transforms.Compose(
+        [
+            transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
+            transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ]
+    )
+
+    def preprocess_train(examples):
+        images = [image.convert("RGB") for image in examples[image_column]]
+        examples["pixel_values"] = [train_transforms(image) for image in images]
+        examples["input_ids"] = tokenize_captions(examples)
+        return examples
+
+    with accelerator.main_process_first():
+        if args.max_train_samples is not None:
+            dataset["train"] = dataset["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
+        # Set the training transforms
+        train_dataset = dataset["train"].with_transform(preprocess_train)
+
+    def collate_fn(examples):
+        pixel_values = torch.stack([example["pixel_values"] for example in examples])
+        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+        input_ids = torch.stack([example["input_ids"] for example in examples])
+        return {"pixel_values": pixel_values, "input_ids": input_ids}
+
+    # DataLoaders creation:
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        shuffle=True,
+        collate_fn=collate_fn,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+    )
+
+    # Prepare everything with our `accelerator`.
+    lora_layers, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        lora_layers, optimizer, train_dataloader, lr_scheduler
+    )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        accelerator.init_trackers("text2image-fine-tune", config=vars(args))
+
+    # Train!
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    global_step = 0
+    first_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+    else:
+        initial_global_step = 0
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=initial_global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+
+    for epoch in range(first_epoch, args.num_train_epochs):
+        unet.train()
+        train_loss = 0.0
+        for step, batch in enumerate(train_dataloader):
+            with accelerator.accumulate(unet):
+                # Convert images to latent space
+                latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
+                latents = latents * vae.config.scaling_factor
+
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents)
+                if args.noise_offset:
+                    # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+                    noise += args.noise_offset * torch.randn(
+                        (latents.shape[0], latents.shape[1], 1, 1), device=latents.device
+                    )
+
+                bsz = latents.shape[0]
+                # Sample a random timestep for each image
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+                timesteps = timesteps.long()
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # Get the text embedding for conditioning
+                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
+
+                # Get the target for loss depending on the prediction type
+                if args.prediction_type is not None:
+                    # set prediction_type of scheduler if defined
+                    noise_scheduler.register_to_config(prediction_type=args.prediction_type)
+
+                if noise_scheduler.config.prediction_type == "epsilon":
+                    target = noise
+                elif noise_scheduler.config.prediction_type == "v_prediction":
+                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
+                else:
+                    raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+                # Predict the noise residual and compute loss
+                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+                if args.snr_gamma is None:
+                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+                else:
+                    # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
+                    # Since we predict the noise instead of x_0, the original formulation is slightly changed.
+                    # This is discussed in Section 4.2 of the same paper.
+                    snr = compute_snr(noise_scheduler, timesteps)
+                    if noise_scheduler.config.prediction_type == "v_prediction":
+                        # Velocity objective requires that we add one to SNR values before we divide by them.
+                        snr = snr + 1
+                    mse_loss_weights = (
+                        torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
+                    )
+
+                    loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+                    loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
+                    loss = loss.mean()
+
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                train_loss += avg_loss.item() / args.gradient_accumulation_steps
+
+                # Backpropagate
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    params_to_clip = lora_layers.parameters()
+                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                accelerator.log({"train_loss": train_loss}, step=global_step)
+                train_loss = 0.0
+
+                if global_step % args.checkpointing_steps == 0:
+                    if accelerator.is_main_process:
+                        # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+                        if args.checkpoints_total_limit is not None:
+                            checkpoints = os.listdir(args.output_dir)
+                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+                            # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+                            if len(checkpoints) >= args.checkpoints_total_limit:
+                                num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+                                removing_checkpoints = checkpoints[0:num_to_remove]
+
+                                logger.info(
+                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+                                )
+                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+                                for removing_checkpoint in removing_checkpoints:
+                                    removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+                                    shutil.rmtree(removing_checkpoint)
+
+                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+                        accelerator.save_state(save_path)
+                        logger.info(f"Saved state to {save_path}")
+
+            logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= args.max_train_steps:
+                break
+
+        if accelerator.is_main_process:
+            if args.validation_prompt is not None and epoch % args.validation_epochs == 0:
+                logger.info(
+                    f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+                    f" {args.validation_prompt}."
+                )
+                # create pipeline
+                pipeline = DiffusionPipeline.from_pretrained(
+                    args.pretrained_model_name_or_path,
+                    unet=accelerator.unwrap_model(unet),
+                    revision=args.revision,
+                    torch_dtype=weight_dtype,
+                )
+                pipeline = pipeline.to(accelerator.device)
+                pipeline.set_progress_bar_config(disable=True)
+
+                # run inference
+                generator = torch.Generator(device=accelerator.device)
+                if args.seed is not None:
+                    generator = generator.manual_seed(args.seed)
+                images = []
+                for _ in range(args.num_validation_images):
+                    images.append(
+                        pipeline(args.validation_prompt, num_inference_steps=30, generator=generator).images[0]
+                    )
+
+                for tracker in accelerator.trackers:
+                    if tracker.name == "tensorboard":
+                        np_images = np.stack([np.asarray(img) for img in images])
+                        tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC")
+                    if tracker.name == "wandb":
+                        tracker.log(
+                            {
+                                "validation": [
+                                    wandb.Image(image, caption=f"{i}: {args.validation_prompt}")
+                                    for i, image in enumerate(images)
+                                ]
+                            }
+                        )
+
+                del pipeline
+                torch.cuda.empty_cache()
+
+    # Save the lora layers
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        unet = unet.to(torch.float32)
+        unet.save_attn_procs(args.output_dir)
+
+        if args.push_to_hub:
+            save_model_card(
+                repo_id,
+                images=images,
+                base_model=args.pretrained_model_name_or_path,
+                dataset_name=args.dataset_name,
+                repo_folder=args.output_dir,
+            )
+            upload_folder(
+                repo_id=repo_id,
+                folder_path=args.output_dir,
+                commit_message="End of training",
+                ignore_patterns=["step_*", "epoch_*"],
+            )
+
+    # Final inference
+    # Load previous pipeline
+    pipeline = DiffusionPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, revision=args.revision, torch_dtype=weight_dtype
+    )
+    pipeline = pipeline.to(accelerator.device)
+
+    # load attention processors
+    pipeline.unet.load_attn_procs(args.output_dir)
+
+    # run inference
+    generator = torch.Generator(device=accelerator.device)
+    if args.seed is not None:
+        generator = generator.manual_seed(args.seed)
+    images = []
+    for _ in range(args.num_validation_images):
+        images.append(pipeline(args.validation_prompt, num_inference_steps=30, generator=generator).images[0])
+
+    if accelerator.is_main_process:
+        for tracker in accelerator.trackers:
+            if len(images) != 0:
+                if tracker.name == "tensorboard":
+                    np_images = np.stack([np.asarray(img) for img in images])
+                    tracker.writer.add_images("test", np_images, epoch, dataformats="NHWC")
+                if tracker.name == "wandb":
+                    tracker.log(
+                        {
+                            "test": [
+                                wandb.Image(image, caption=f"{i}: {args.validation_prompt}")
+                                for i, image in enumerate(images)
+                            ]
+                        }
+                    )
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()