Upload DATA/experiment.py with huggingface_hub

DATA/experiment.py

@@ -0,0 +1,110 @@
+import os
+from PIL import Image
+import torchvision.transforms.functional as f
+from utils import load_face_generator
+from omegaconf import OmegaConf
+
+def generate_face_image(
+    anything_facemaker,
+    class_concept, 
+    face_img_pil=None, 
+    controlnet_conditioning_scale=1.0,
+    strength=0.95,
+):
+    face_img_pil = f.center_crop(
+        f.resize(face_img_pil, 512), 512).convert('RGB')
+    prompt = anything_facemaker.prompt_template.format(class_concept)
+    # # There are four ways to generate a image by now.
+    # pure_generate = anything_facemaker.generate(prompt=prompt, image=face_img_pil, do_inversion=False)
+    # inversion = anything_facemaker.generate(prompt=prompt, image=face_img_pil, strength=strength, do_inversion=True)
+
+    if controlnet_conditioning_scale == None:
+        init_face_pil = anything_facemaker.generate(prompt=prompt)
+        return init_face_pil
+    
+    if strength is None:
+        pure_control = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=False,
+                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale)
+        init_face_pil = pure_control
+    else:
+        control_inversion = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=True, 
+                                                                 strength=strength,
+                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale)
+        init_face_pil = control_inversion
+    return init_face_pil
+
+
+def experiment(anything_facemaker, concepts_path, face_img_path, output_dir,
+               controlnet_conditioning_scale=1., strength=0.95):
+    os.makedirs(output_dir, exist_ok=True)
+    face_img_pil = Image.open(face_img_path)
+    face_img_pil = f.center_crop(
+        f.resize(face_img_pil, 512), 512).convert('RGB')
+    with open(concepts_path) as fr:
+        concepts = fr.read().split('\n')
+    concepts = [concept for concept in concepts if len(concept)!=0]
+    for concept in concepts:
+        save_path = os.path.join(output_dir, f'{concept}.png')
+        if os.path.exists(save_path):
+            continue
+        init_face_pil = generate_face_image(
+            anything_facemaker,
+            class_concept=concept,
+            face_img_pil=face_img_pil,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            strength=strength,
+        )
+
+        save_path = os.path.join(output_dir, f'{concept}.png')
+        init_face_pil.save(save_path)
+
+
+
+if __name__=='__main__':
+    # run this in repo path:
+    # PYTHONPATH=.:$PYTHONPATH python experiments/experiment.py 
+    model_config_path = 'resources/models.yaml'
+    # model_config_path = 'resources/models_personality.yaml'
+    model_config = OmegaConf.load(model_config_path)['models']
+    gameicon_config = model_config['GameIconInstitute_mode']
+
+    face_img_path='resources/images/faces/1.jpg'
+    controlnet_conditioning_scale=1.
+    strength=0.95
+
+    for model, model_info in model_config.items():
+        
+        anything_facemaker = load_face_generator(
+            model_dir=model_info['model_dir'],
+            lora_path=model_info['lora_path'],
+            prompt_template=model_info['prompt_template'],
+            negative_prompt=model_info['negative_prompt']
+        )
+        output_dir = os.path.join('results/concepts_test/control_inversion', model)
+        os.makedirs(output_dir, exist_ok=True)
+        # concept test, with control and inversion
+        input_dir = 'resources/prompts'
+        for dir, folders, files in os.walk(input_dir):
+            for file in files:
+                input_file = os.path.join(dir, file)
+                file_output_dir = os.path.join(output_dir, file)
+                print(f'input_file: {input_file}')
+                print(f'file_output_dir: {file_output_dir}')
+                experiment(anything_facemaker, input_file, face_img_path, output_dir=file_output_dir,
+                        controlnet_conditioning_scale=controlnet_conditioning_scale,
+                        strength=strength)
+
+    # # concept, with control and inversion
+    # experiment(anything_facemaker, 'resources/concepts.txt', face_img_path, output_dir='results/concepts/control_inversion',
+    #         controlnet_conditioning_scale=controlnet_conditioning_scale,
+    #         strength=strength)
+
+    # # concept test, no control no inversion
+    # experiment(anything_facemaker, 'resources/concepts_test.txt', face_img_path, output_dir='results/concepts_test/generate',
+    #         controlnet_conditioning_scale=None,
+    #         strength=strength)
+
+    # # concept, no control no inversion
+    # experiment(anything_facemaker, 'resources/concepts.txt', face_img_path, output_dir='results/concepts/generate',
+    #         controlnet_conditioning_scale=None,
+    #         strength=strength)