reproduce.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xulu/anaconda3/envs/diffusion2/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from diffusers import StableDiffusionPipeline, UNet2DConditionModel\n",
    "import torch\n",
    "import json\n",
    "from utils.misc import *\n",
    "import gc\n",
    "import pandas as pd\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from utils.clip_eval import CLIPEvaluator\n",
    "from transformers import AutoImageProcessor, AutoModel\n",
    "from utils.layout_control import *\n",
    "from transformers import CLIPTextModel\n",
    "from matplotlib import pyplot as plt\n",
    "import math\n",
    "\n",
    "\n",
    "BACKGROUND = ['barn']\n",
    "OBJ = ['cat', 'dog', 'chair', 'table', 'flower', 'wooden_pot']\n",
    "\n",
    "\n",
    "def init_generative_model(args, device):\n",
    "    \"\"\"\n",
    "    Initialize the model\n",
    "    params:\n",
    "        args: argparse.Namespace\n",
    "        device: str\n",
    "    \"\"\"\n",
    "    if args.model_name == \"sd-v1-5\":\n",
    "        model_id = \"runwayml/stable-diffusion-v1-5\"\n",
    "        pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)\n",
    "    elif args.model_name == 'textual_inversion':\n",
    "        model_id = \"runwayml/stable-diffusion-v1-5\"\n",
    "        if args.spatial_inversion:\n",
    "            pipe = Layout_Control(model_id=model_id, device=device)\n",
    "        else:\n",
    "            pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)\n",
    "    elif args.model_name == 'custom_diffusion' or args.model_name == 'dreambooth':\n",
    "        model_id = \"CompVis/stable-diffusion-v1-4\"\n",
    "        if args.spatial_inversion:\n",
    "            pipe = Layout_Control(model_id=model_id, device=device)\n",
    "        else:\n",
    "            pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)\n",
    "    return pipe\n",
    "\n",
    "\n",
    "def check_identifier_token_dir(checkpoint, concept_str):\n",
    "    if os.path.exists(os.path.join(checkpoint, concept_str)):\n",
    "        dir = os.path.join(checkpoint, concept_str)\n",
    "    elif os.path.exists(checkpoint.format(concept_str)):\n",
    "        dir = checkpoint.format(concept_str)\n",
    "    else:\n",
    "        dir = checkpoint\n",
    "    return dir\n",
    "\n",
    "\n",
    "def load_trained_weights(args, pipe, checkpoint, concepts_str, c_identifier):\n",
    "    \"\"\"\n",
    "    Load additional weights for the model\n",
    "    params:\n",
    "        model_name: str\n",
    "        pipe: StableDiffusionPipeline\n",
    "        checkpoint: str, where the additional unique identifier weights are saved\n",
    "        concepts_str: str, the concepts used to generate images\n",
    "    \"\"\"\n",
    "    print(f\"loading unique identifier weights for {concepts_str} ...\")\n",
    "    if args.model_name == 'sd-v1-5':\n",
    "        return pipe\n",
    "    if args.model_name == 'textual_inversion':\n",
    "        for cls in concepts_str.split(','):\n",
    "            dir = check_identifier_token_dir(checkpoint, cls)\n",
    "            pipe.load_textual_inversion(os.path.join(dir, 'learned_embeds.bin'))\n",
    "    elif args.model_name == 'custom_diffusion' and args.spatial_inversion == False:\n",
    "        dir = check_identifier_token_dir(checkpoint, concepts_str)\n",
    "        print(f\"loading unique identifier weights from {dir} ...\")\n",
    "        pipe.unet.load_attn_procs(dir, weight_name=\"pytorch_custom_diffusion_weights.bin\")\n",
    "        for cls in concepts_str.split(','):\n",
    "            pipe.load_textual_inversion(dir, weight_name=f\"{c_identifier[cls]}.bin\")\n",
    "    elif args.model_name == 'custom_diffusion' and args.spatial_inversion == True:\n",
    "        dir = check_identifier_token_dir(checkpoint, concepts_str)\n",
    "        pipe.load_attn_procs(dir, weight_name=\"pytorch_custom_diffusion_weights.bin\")\n",
    "        for cls in concepts_str.split(','):\n",
    "            pipe.load_textual_inversion(os.path.join(dir, f\"{c_identifier[cls]}.bin\"))\n",
    "    elif args.model_name == 'dreambooth' and args.spatial_inversion == False:\n",
    "        dir = check_identifier_token_dir(checkpoint, concepts_str)\n",
    "        unet = UNet2DConditionModel.from_pretrained(os.path.join(dir, 'unet')).to(pipe.device)\n",
    "        pipe.unet = unet\n",
    "        if os.path.exists(os.path.join(dir, 'text_encoder')):\n",
    "            text_encoder = CLIPTextModel.from_pretrained(os.path.join(dir, 'text_encoder')).to(pipe.device)\n",
    "            pipe.text_encoer = text_encoder\n",
    "        for cls in concepts_str.split(','):\n",
    "            if os.path.exists(os.path.join(dir, f\"{c_identifier[cls]}.bin\")):\n",
    "                print(f\"loading unique identifier weights from {c_identifier[cls]}.bin ...\")\n",
    "                pipe.load_textual_inversion(os.path.join(dir, f\"{c_identifier[cls]}.bin\"))\n",
    "    elif args.model_name == 'dreambooth' and args.spatial_inversion == True:\n",
    "        dir = check_identifier_token_dir(checkpoint, concepts_str)\n",
    "        pipe.load_dreambooth_weights(dir)\n",
    "        for cls in concepts_str.split(','):\n",
    "            pipe.load_textual_inversion(os.path.join(dir, f\"{c_identifier[cls]}.bin\"))\n",
    "    else:\n",
    "        raise ValueError(f'Unknown model name {args.model_name}')\n",
    "    return pipe\n",
    "\n",
    "\n",
    "def save_img(im, prompts, save_dir, name):\n",
    "    \"\"\"\n",
    "    Save images\n",
    "    params:\n",
    "        im: numpy array\n",
    "        prompts: list\n",
    "        save_dir: str\n",
    "        name: str\n",
    "    \"\"\"\n",
    "    save_path = os.path.join(save_dir, f'{prompts}/{name}.jpg')\n",
    "    check_mk_file_dir(save_path)\n",
    "    if isinstance(im, np.ndarray):\n",
    "        Image.fromarray(im).save(save_path)\n",
    "    elif isinstance(im, Image.Image):\n",
    "        im.save(save_path)\n",
    "    else:\n",
    "        raise TypeError(f'Unknown type {type(im)}') \n",
    "\n",
    "def get_reference_images(concepts_str, src_img_dir):\n",
    "    src_imgs = {}\n",
    "    for concept in concepts_str.split(','):\n",
    "        src_imgs[concept] = []\n",
    "        for img_path in os.listdir(os.path.join(src_img_dir, concept)):\n",
    "            src_imgs[concept].append(Image.open(os.path.join(src_img_dir, concept, img_path)))\n",
    "    return src_imgs\n",
    "\n",
    "\n",
    "def edit_original_prompt(prompt, c_identifier, keys, mode='replace'):\n",
    "    \"\"\"\n",
    "    Edit the original prompt to the prompt with identifiers\n",
    "    params:\n",
    "        prompt: str\n",
    "        c_identifier: dict, concepts and identifiers, e.g. {'cat': '<cute-cat>'}\n",
    "        keys: concept list, e.g. ['cat', 'dog']\n",
    "        mode: str, 'replace' or 'insert'\n",
    "    returns:\n",
    "        prompt: str\n",
    "    \"\"\"\n",
    "    if mode != 'none':\n",
    "        for key in keys:\n",
    "            if '_' in key:\n",
    "                replace_str = key.replace('_', ' ')\n",
    "            else:\n",
    "                replace_str = key\n",
    "            if replace_str in prompt:\n",
    "                if mode == 'replace':\n",
    "                    prompt = prompt.replace(replace_str, c_identifier[key])\n",
    "                elif mode == 'insert':\n",
    "                    prompt = prompt.replace(replace_str, f'{c_identifier[key]} {replace_str}')\n",
    "            else:\n",
    "                raise ValueError(f'{replace_str} not in prompt {prompt}')\n",
    "    return prompt.strip()\n",
    "\n",
    "\n",
    "def get_input_dict(prompt, edited_prompt, concepts_str, c_identifier):\n",
    "    \"\"\"\n",
    "    Get the input dict for the layout control generation\n",
    "    params:\n",
    "        prompt: str\n",
    "        edited_prompt: str\n",
    "        concepts_str: str\n",
    "        c_identifier: dict, concepts and identifiers, e.g. {'cat': '<cute-cat>'}\n",
    "        layout: dict, {prompt: [bbox]}, e.g. {'prompt': [[[0.3, 0.4, 0.5, 0.7]], [[0.5, 0.4, 0.7, 0.8]]]}\n",
    "    \"\"\"\n",
    "    phrase_list = []\n",
    "    for c in concepts_str.split(','):\n",
    "        if c in prompt:\n",
    "            phrase_list.append(c)\n",
    "        elif c.replace('_', ' ') in prompt:\n",
    "            phrase_list.append(c.split('_')[-1])\n",
    "    phrases = '; '.join(phrase_list)\n",
    "    layout_info = get_constant_layout(prompt, concepts_str)\n",
    "    identifier = layout_info['objs']\n",
    "    for c in concepts_str.split(','):\n",
    "        identifier = identifier.replace(c, c_identifier[c])\n",
    "    input_dict = {\"prompt\": prompt,\n",
    "                \"edited_prompt\": edited_prompt,\n",
    "                \"phrases\": phrases,\n",
    "                \"identifier\": identifier,\n",
    "                \"bboxes\": layout_info['bboxes']}\n",
    "    print(input_dict)\n",
    "    return input_dict\n",
    "\n",
    "\n",
    "def get_constant_layout(prompt, concepts_str):\n",
    "    if ' and ' in prompt:\n",
    "        obj1 = prompt.split(' and ')[0].strip()\n",
    "        obj2 = prompt.split(' and ')[-1].strip()\n",
    "    elif 'in front of' in prompt:\n",
    "        obj1 = prompt.split(' in front of ')[-1].strip()\n",
    "        obj2 = prompt.split(' in front of ')[0].strip()\n",
    "    else:\n",
    "        obj1, obj2 = concepts_str.split(',')\n",
    "    objs = []\n",
    "    bboxes = []\n",
    "    concepts = concepts_str.split(',')\n",
    "    if len(concepts) == 1:\n",
    "        objs.extend(concepts)\n",
    "        bboxes.append([[0.1, 0.2, 0.5, 0.8]])\n",
    "        for word in obj2.split(' '):\n",
    "            objs.append(word)\n",
    "            bboxes.append([[0.6, 0.2, 0.95, 0.8]])\n",
    "    elif len(concepts) == 2:\n",
    "        objs.append(obj1.replace(' ', '_'))\n",
    "        bboxes.append([[0.1, 0.2, 0.5, 0.8]])\n",
    "        objs.append(obj2.replace(' ', '_'))\n",
    "        bboxes.append([[0.6, 0.2, 0.95, 0.8]])\n",
    "    objs = '; '.join(objs)\n",
    "    return {'objs': objs, 'bboxes': bboxes}\n",
    "\n",
    "\n",
    "def generate_images(args, pipe, c_p, c_identifier):\n",
    "    \"\"\"\n",
    "    Generate images from the optimization-based model\n",
    "    params:\n",
    "        args: argparse.Namespace\n",
    "        pipe: StableDiffusionPipeline\n",
    "        c_p: dict, concepts and prompts, e.g. {'cat': ['a photo of a cat']}\n",
    "        c_identifier: dict, concepts and identifiers, e.g. {'cat': '<cute-cat>'}\n",
    "    \"\"\"\n",
    "    for concepts_str, prompts in c_p.items():\n",
    "        try:\n",
    "            pipe = load_trained_weights(args, pipe, args.checkpoint, concepts_str, c_identifier)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(\"Loading unique identifier failed, re-initializing the model...\")\n",
    "            del pipe\n",
    "            gc.collect()\n",
    "            torch.cuda.empty_cache()\n",
    "            pipe = init_generative_model(args, args.device)\n",
    "            pipe = load_trained_weights(args, pipe, args.checkpoint, concepts_str, c_identifier)\n",
    "\n",
    "        for prompt in prompts:\n",
    "            edited_prompt = edit_original_prompt(prompt, c_identifier, concepts_str.split(','), mode=args.edit_mode)\n",
    "            print(\"Generating images for prompt: \", edited_prompt)\n",
    "            generator = torch.manual_seed(8888)\n",
    "            if args.spatial_inversion:\n",
    "                input_dict = get_input_dict(prompt, edited_prompt, concepts_str, c_identifier)\n",
    "                for idx in range(args.num_per_prompt):\n",
    "                    im = pipe(input_dict, num_steps=50, guidance_scale=7.5, generator=generator)[0]\n",
    "                    save_img(im, prompt, os.path.join(args.img_save_dir, concepts_str), idx)\n",
    "            else:\n",
    "                for idx in range(args.num_per_prompt):\n",
    "                    im = pipe(edited_prompt, num_inference_steps=50, guidance_scale=7.5, generator=generator).images[0]\n",
    "                    save_img(im, prompt, os.path.join(args.img_save_dir, concepts_str), idx)\n",
    "    print(\"all images saved in \", args.img_save_dir)\n",
    "\n",
    "\n",
    "def check_dir(img_save_dir, concepts_str, prompt):\n",
    "    \"\"\"\n",
    "    check dir if exists for alignment function\n",
    "    \"\"\"\n",
    "    if os.path.exists(os.path.join(img_save_dir, concepts_str, prompt)):\n",
    "        dir = os.path.join(img_save_dir, concepts_str, prompt)\n",
    "    else:\n",
    "        raise ValueError(f'no such dir: {os.path.join(img_save_dir, concepts_str, prompt)}')\n",
    "    return dir\n",
    "\n",
    "\n",
    "def eval_alignment(evaluator, c_p, img_save_dir, src_img_dir):\n",
    "    \"\"\"\n",
    "    Evaluate the alignment between the generated images and the source images/texts\n",
    "    params:\n",
    "        concepts_list: list\n",
    "        evaluator: CLIPEvaluator\n",
    "        c_p: dict, concepts and prompts\n",
    "        img_save_dir: str, where the generated images are saved\n",
    "        src_img_dir: str, where the source images are saved\n",
    "    returns:\n",
    "        img_img_sim_mean: dict, the average similarity between the generated images and the source images\n",
    "        text_img_sim_mean: dict, the average similarity between the generated images and the source texts\n",
    "    \"\"\"\n",
    "    img_img_sim = {}   # {c_str: {c1:[s1, s2, s3], c2:[s1, s2, s3]}}\n",
    "    text_img_sim = {}  # {c_str: [s1, s2, s3]}\n",
    "    \n",
    "    for concepts_str, prompts in c_p.items():\n",
    "        print(f'evaluating {concepts_str}...')\n",
    "        img_img_sim[concepts_str] = {} \n",
    "        text_img_sim[concepts_str] = [] \n",
    "\n",
    "        src_imgs = get_reference_images(concepts_str, src_img_dir) # {c1: [I1, I2...], c2: [I1, I2...]}\n",
    "        src_img_features = {}  # {c1: [f1, f2, f3...], c2: [f1, f2, f3...}\n",
    "        for concept, concept_src_imgs in src_imgs.items():\n",
    "            src_img_features[concept] = [evaluator.get_image_features(src_img) for src_img in concept_src_imgs]\n",
    "            img_img_sim[concepts_str][concept] = []\n",
    "\n",
    "        for prompt in prompts:\n",
    "            text_feature = evaluator.get_text_features(prompt)\n",
    "            dir = check_dir(img_save_dir, concepts_str, prompt)\n",
    "            for img_path in os.listdir(dir):\n",
    "                img = Image.open(os.path.join(dir, img_path))\n",
    "                # text alignment\n",
    "                text_img_sim[concepts_str].append(2.5*evaluator.txt_to_img_similarity(img, text_features=text_feature).cpu().numpy())\n",
    "                # image alignment\n",
    "                for concept, concept_img_features in src_img_features.items():\n",
    "                    img_img_sim[concepts_str][concept].extend([evaluator.img_to_img_similarity(img, src_img_features=src_img_feature).cpu().numpy() for src_img_feature in concept_img_features])\n",
    "    \n",
    "    img_img_sim_mean = {}\n",
    "    text_img_sim_mean = {}\n",
    "    for concepts_str, sim_list in img_img_sim.items():\n",
    "        img_img_sim_mean[concepts_str] = {}\n",
    "        for concept, sims in sim_list.items():\n",
    "            img_img_sim_mean[concepts_str][concept] = np.mean(sims)\n",
    "        text_img_sim_mean[concepts_str] = np.mean(text_img_sim[concepts_str])\n",
    "        print(f'{concepts_str}: image alignment -- {img_img_sim_mean[concepts_str]} | text alignment -- {text_img_sim_mean[concepts_str]}')\n",
    "    return {'img': img_img_sim_mean, 'text': text_img_sim_mean}\n",
    "\n",
    "def eval_coco_objs(evaluator, c_p, img_save_dir):\n",
    "    \"\"\"\n",
    "    Evaluate the alignment between the generated images and the source images/texts\n",
    "    params:\n",
    "        evaluator: CLIPEvaluator\n",
    "        c_p: dict, concepts and prompts\n",
    "        img_save_dir: str, where the generated images are saved\n",
    "    returns:\n",
    "        coco_coi_mean: dict, the average similarity between the generated images and the source images\n",
    "    \"\"\"\n",
    "    coco_objs = {}\n",
    "    for concepts_str, prompts in c_p.items():\n",
    "        print(f'evaluating {concepts_str} coco objects...')\n",
    "        coco_objs[concepts_str] = []\n",
    "        for prompt in prompts:\n",
    "            if os.path.exists(os.path.join(img_save_dir, concepts_str, prompt)):\n",
    "                dir = os.path.join(img_save_dir, concepts_str, prompt)\n",
    "            else:\n",
    "                raise ValueError(f'no such dir: {os.path.join(img_save_dir, concepts_str, prompt)}')\n",
    "            for img_path in os.listdir(dir):\n",
    "                img = Image.open(os.path.join(dir, img_path))\n",
    "                acc = evaluator.detect_objects(img, prompt)\n",
    "                coco_objs[concepts_str].append(acc)\n",
    "    coco_objs_coi = {}\n",
    "    for concepts_str in coco_objs.keys():\n",
    "        coco_objs_coi[concepts_str] = np.mean(coco_objs[concepts_str])\n",
    "        print(f'{concepts_str}: coco CoI -- {coco_objs_coi[concepts_str]}')\n",
    "    return {'coco_coi': coco_objs_coi}\n",
    "\n",
    "\n",
    "def get_placeholders(concepts_list):\n",
    "    \"\"\"\n",
    "    Get the placeholders for the concepts\n",
    "    params:\n",
    "        concepts_list: list, the concepts list\n",
    "    \"\"\"\n",
    "    cls_identifier = {}\n",
    "    for concept in concepts_list:\n",
    "        cls_identifier[concept['class_prompt']] = concept['placeholder']\n",
    "    return cls_identifier\n",
    "\n",
    "\n",
    "def get_concept_prompts(concepts_list, num_inverted_concepts):\n",
    "    \"\"\"\n",
    "    Get the concept prompts\n",
    "    params:\n",
    "        concepts_list: list, the concepts list\n",
    "        num_inverted_concepts: int, the number of inverted concepts\n",
    "    returns:\n",
    "        c_p: dict, {concept_str: [prompts]}, e.g. {'cat_dog': ['a cat in front of a dog', 'a dog playing with a cat']}\n",
    "    \"\"\"\n",
    "    if num_inverted_concepts == 1:\n",
    "        concepts = [concept['class_prompt'] for concept in concepts_list]\n",
    "        c_p = {}\n",
    "        with open(\"./data/coco.txt\", 'r') as f:\n",
    "            coco_objs = f.read().split('\\n')\n",
    "        for concept in concepts:\n",
    "            conj = 'in front of' if concept in BACKGROUND else 'and'\n",
    "            c_p[concept] = []\n",
    "            for obj in coco_objs:\n",
    "                obj = obj.strip()\n",
    "                if conj == 'and':\n",
    "                    c_p[concept].append(f\"{concept.replace('_', ' ')} {conj} {obj}\".strip())\n",
    "                else:\n",
    "                    c_p[concept].append(f\"{obj} {conj} {concept.replace('_', ' ')}\".strip())      \n",
    "        return c_p\n",
    "    if num_inverted_concepts == 2:\n",
    "        c_p = {}\n",
    "        concepts = [concept['class_prompt'] for concept in concepts_list]\n",
    "        for i in range(len(concepts)):\n",
    "            for j in range(i+1, len(concepts)):\n",
    "                conj = 'and'\n",
    "                c_p[f'{concepts[i]},{concepts[j]}'] = []\n",
    "                c_p[f'{concepts[i]},{concepts[j]}'].append(f\"{concepts[i].replace('_', ' ')} {conj} {concepts[j].replace('_', ' ')}\")\n",
    "                c_p[f'{concepts[i]},{concepts[j]}'].append(f\"{concepts[j].replace('_', ' ')} {conj} {concepts[i].replace('_', ' ')}\")\n",
    "    return c_p\n",
    "\n",
    "def generate(args):\n",
    "    set_seed(seed=8888)\n",
    "    # prepare the prompts\n",
    "    concepts_list = json.load(open(args.concepts_list_path, \"r\"))\n",
    "    c_p = get_concept_prompts(concepts_list, args.num_inverted_concepts)   # dict: {concept_str: [prompts]}\n",
    "    c_identifier = get_placeholders(concepts_list)\n",
    "    # prepare the geneartive model\n",
    "    pipe = init_generative_model(args, args.device)\n",
    "    # generate images\n",
    "    generate_images(args, pipe, c_p, c_identifier)\n",
    "\n",
    "\n",
    "def evaluate(args, eval_clip_score=True, eval_coco_coi=False):\n",
    "    \"\"\"\n",
    "    Evaluate the generated images\n",
    "    params:\n",
    "        args: argparse.Namespace\n",
    "        eval_clip_score: bool, whether to evaluate the CLIP score (including image alignment and text-image alignment)\n",
    "        eval_coco_coi: bool, whether to evaluate the coco CoI score\n",
    "    \"\"\"\n",
    "    results = {}\n",
    "    concepts_list = json.load(open(args.concepts_list_path, \"r\"))\n",
    "    c_p = get_concept_prompts(concepts_list, args.num_inverted_concepts)   # dict: {concept_str: [prompts]}\n",
    "    \n",
    "    if eval_clip_score:\n",
    "        # evaluate the image alignment and text-image alignment\n",
    "        evaluator = CLIPEvaluator(args.device)\n",
    "        clip_score = eval_alignment(evaluator, c_p, args.img_save_dir, args.src_img_dir)\n",
    "        results.update(clip_score)\n",
    "\n",
    "    if args.num_inverted_concepts==1 and eval_coco_coi==True:\n",
    "        coco_coi = eval_coco_objs(evaluator, c_p, args.img_save_dir)\n",
    "        results.update(coco_coi)\n",
    "        \n",
    "    #KID = eval_KID(ims)\n",
    "    save_results(args, results)\n",
    "\n",
    "def run_and_test(args, **kwargs):\n",
    "    generate(args)\n",
    "    evaluate(args, **kwargs)\n",
    "    \n",
    "def save_results(args, scores):\n",
    "    \"\"\" \n",
    "    Save the results to a csv file\n",
    "    params:\n",
    "        args: argparse.Namespace\n",
    "        scores: dict, {score_name: {concept_str: score}}\n",
    "    \"\"\"\n",
    "    results = args.get_dict()\n",
    "    if os.path.exists(args.results_path):\n",
    "        file = pd.read_csv(args.results_path)\n",
    "    else:\n",
    "        file = pd.DataFrame(columns=list(results.keys()))\n",
    "    \n",
    "    for score_name, score_dict in scores.items():\n",
    "        for concept_str, score in score_dict.items():\n",
    "            if isinstance(score, dict):\n",
    "                for concept, s in score.items():\n",
    "                    results[f'{concept_str}_{concept}_{score_name}'] = s\n",
    "            else:\n",
    "                results[f'{concept_str}_{score_name}'] = score\n",
    "    file = pd.concat([file, pd.DataFrame(results, index=[0])], ignore_index=True)\n",
    "    file.to_csv(args.results_path, index=False)\n",
    "    print(\"results saved in \", args.results_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {\n",
    "    \"model_name\": \"custom_diffusion\",  # textual_inversion, dreambooth, custom_diffusion\n",
    "    \"spatial_inversion\": True,   # True or False\n",
    "    \"device\": \"cuda:0\",\n",
    "    \"num_inverted_concepts\": 1,   # 1 or 2\n",
    "    \"edit_mode\": \"insert\",\n",
    "    \"num_per_prompt\": 5,  # number of generated images per prompt\n",
    "    \"concepts_list_path\": \"./data/concepts_list_test.json\",\n",
    "    \"checkpoint\": \"./snapshot/compositional_custom_diffusion/{}\",   \n",
    "    \"src_img_dir\": \"./data/reference_images/\",\n",
    "    \"results_path\": \"results.csv\",\n",
    "}\n",
    "args[\"img_save_dir\"] = \"./samples/{}_spatial_{}/{}c\".format(args['model_name'],args['spatial_inversion'], args['num_inverted_concepts'])\n",
    "args = Dict2Class(args)\n",
    "run_and_test(args, eval_clip_score=True, eval_coco_coi=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "diffusion2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}