update code

app.py

@@ -271,6 +271,8 @@ def prepare_working_dir(dir, style):
 def launch_pretrained():
     from huggingface_hub import hf_hub_download, snapshot_download
     hf_hub_download(repo_id="KumaPower/AvatarArtist", repo_type='model', local_dir="./pretrained_model")
+    hf_hub_download(repo_id="stabilityai/stable-diffusion-2-1-base", repo_type='model', local_dir="./pretrained_model/sd21")
+    hf_hub_download(repo_id="CrucibleAI/ControlNetMediaPipeFace", repo_type='model', local_dir="./pretrained_model/control")
 
 
 def prepare_image_list(img_dir, selected_img):
@@ -394,8 +396,8 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
     label_file_test = os.path.join(target_path, 'images512x512/dataset_realcam.json')
 
     if source_type == 'example':
-        input_img_fvid = './demo_data/source_img/img_generate_different_domain/coeffs/trained_input_imgs'
-        input_img_motion = './demo_data/source_img/img_generate_different_domain/motions/trained_input_imgs'
+        input_img_fvid = './demo_data/source_img/img_generate_different_domain/coeffs/demo_imgs'
+        input_img_motion = './demo_data/source_img/img_generate_different_domain/motions/demo_imgs'
     elif source_type == 'custom':
         input_img_fvid = os.path.join(save_path_base, 'processed_img/dataset/coeffs/input_image')
         input_img_motion = os.path.join(save_path_base, 'processed_img/dataset/motions/input_image')
@@ -704,7 +706,7 @@ def launch_gradio_app():
             """
         )
         # DISPLAY
-        image_folder = "./demo_data/source_img/img_generate_different_domain/images512x512/trained_input_imgs"
+        image_folder = "./demo_data/source_img/img_generate_different_domain/images512x512/demo_imgs"
         video_folder = "./demo_data/target_video"
 
         examples_images = sorted(
@@ -886,17 +888,17 @@ def launch_gradio_app():
 
 if __name__ == '__main__':
     import torch.multiprocessing as mp
-
     mp.set_start_method('spawn', force=True)
-    image_folder = "./demo_data/source_img/img_generate_different_domain/images512x512/trained_input_imgs"
+    launch_pretrained()
+    image_folder = "./demo_data/source_img/img_generate_different_domain/images512x512/demo_imgs"
     example_img_names = os.listdir(image_folder)
     render_model, sample_steps, DiT_model, \
         vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, std, mean, ws_avg, Faceverse, device, input_process_model = model_define()
-    controlnet_path = '/nas8/liuhongyu/model/ControlNetMediaPipeFaceold'
+    controlnet_path = './pretrained_model/control'
     controlnet = ControlNetModel.from_pretrained(
         controlnet_path, torch_dtype=torch.float16
     )
-    sd_path =  '/nas8/liuhongyu/model/stable-diffusion-2-1-base'
+    sd_path =  './pretrained_model/sd21'
     pipeline_sd = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
         sd_path, torch_dtype=torch.float16,
         use_safetensors=True, controlnet=controlnet, variant="fp16"