Upload 6 files

README.md

@@ -1,14 +1,51 @@
----
-title: Stablehairv2 Demo
-emoji: 📈
-colorFrom: pink
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.44.1
-app_file: app.py
-pinned: false
-license: mit
-short_description: This is a simple demo showing our Huawei cup work
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# StableHair v2
+**Stable-Hair v2: Real-World Hair Transfer via Multiple-View Diffusion Model**  
+Kuiyuan Sun*, [Yuxuan Zhang*](https://xiaojiu-z.github.io/YuxuanZhang.github.io/), [Jichao Zhang*](https://zhangqianhui.github.io/), [Jiaming Liu](https://scholar.google.com/citations?user=SmL7oMQAAAAJ&hl=en), 
+ [Wei Wang](https://weiwangtrento.github.io/), [Nicu Sebe](http://disi.unitn.it/~sebe/), [Yao Zhao](https://scholar.google.com/citations?user=474TbQYAAAAJ&hl=en&oi=ao)<br>
+*Equal Contribution <br>
+Beijing Jiaotong University, Shanghai Jiaotong University, Ocean University of China, Tiamat AI, University of Trento <br>
+[Arxiv](https://arxiv.org/abs/2507.07591), [Project]()<br>
+
+
+Bald     |  Reference | Multiple View | Original Video
+![](./imgs/multiview1.gif)  
+Bald     |  Reference | Multiple View | Original Video
+![](./imgs/multiview2.gif)
+
+## Environments
+
+```
+conda create -n stablehairv2 python=3.10
+```
+```
+pip install -r requirements.txt
+```
+
+## Results
+
+<img src="./imgs/teaser.jpg" width="800"> 
+
+## Pretrained Model
+| Name                       |   Model   | 
+|----------------------------|:---------:|
+| motion_module-41400000.pth | [:link:](https://drive.google.com/file/d/1AZMhui9jNRF3Z0N72VDPOwDd0JafLQ3B/view?usp=drive_link) | 
+| pytorch_model_1.bin        |  [:link:](https://drive.google.com/file/d/1FwKPZI8lvdlZqu8R1aJ-QbE55kxHPHjU/view?usp=drive_link) | 
+| pytorch_model_2.bin        |  [:link:](https://drive.google.com/file/d/1h3dXlo8lhZN3ee5aN0shZmpLfn5itVou/view?usp=drive_link) | 
+| pytorch_model_3.bin        |  [:link:](https://drive.google.com/file/d/1jARfXaU6wiur85Vm1JxZ_xye0FfrUiqb/view?usp=drive_link) | 
+| pytorch_model.bin          |  [:link:](https://drive.google.com/file/d/1zXXf13pV5IOn2vrV6DGI9hliEFvuPrYf/view?usp=drive_link) |
+
+### Multiple View Hair Transfer
+
+Please use ``gdown''' to download the pretrained model and save it in your model_path
+```
+
+python test_stablehairv2.py --pretrained_model_name_or_path "stable-diffusion-v1-5/stable-diffusion-v1-5" \
+       --image_encoder "openai/clip-vit-large-patch14" --output_dir [Your_output_dir]  \
+       --num_validation_images 1 --validation_ids ./test_imgs/bald.jpg \
+       --validation_hairs ./test_imgs/ref1.jpg --model_path [Your_model_path]
+```
+
+
+# Our V1 version
+
+StableHair v2 is an improved version of [StableHair](https://github.com/Xiaojiu-z/Stable-Hair) (AAAI 2025)

dataset_mv.py

@@ -0,0 +1,2236 @@
+from torch.utils import data
+import os
+import torch
+import numpy as np
+import cv2
+import random
+import albumentations as A
+
+pixel_transform = A.Compose([
+    A.SmallestMaxSize(max_size=512),
+    A.CenterCrop(512, 512),
+    A.Affine(scale=(0.5, 1), translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, rotate=(-10, 10), p=0.8),
+], additional_targets={'image0': 'image', 'image1': 'image'})
+
+hair_transform = A.Compose([
+    A.SmallestMaxSize(max_size=512),
+    A.CenterCrop(512, 512),
+    A.Affine(scale=(0.9, 1.2), rotate=(-10, 10), p=0.7)]
+)
+
+# class myDataset(data.Dataset):
+#     """Custom data.Dataset compatible with data.DataLoader."""
+
+#     def __init__(self, train_data_dir):
+#         self.img_path = os.path.join(train_data_dir, "hair")
+#         # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+#         # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+#         # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+#         self.pose_path = os.path.join(train_data_dir, "pose.npy")
+#         self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+#         self.ref_path = os.path.join(train_data_dir, "reference")
+#         self.lists = os.listdir(self.img_path)
+#         self.len = len(self.lists)-10
+#         self.pose = np.load(self.pose_path)
+#         #self.pose = np.random.randn(12, 4)
+
+#     def __getitem__(self, index):
+#         """Returns one data pair (source and target)."""
+#         # seq_len, fea_dim
+#         random_number1 = random.randrange(0, 21)
+#         random_number2 = random.randrange(0, 21)
+
+#         while random_number2 == random_number1:
+#             random_number2 = random.randrange(0, 21)
+#         name = self.lists[index]
+
+#         random_number1 = random_number1
+#         #* 10
+#         #random_number2 = random_number2 * 10
+
+#         random_number2 = random_number1
+
+        
+#         non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+#         ref_folder = os.path.join(self.ref_path, name)
+
+#         files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+#         ref_path = os.path.join(ref_folder, files[0])
+        
+#         img_non_hair = cv2.imread(non_hair_path)
+#         ref_hair = cv2.imread(ref_path)
+
+        
+#         img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+#         ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+#         img_non_hair = cv2.resize(img_non_hair, (512, 512))
+#         ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+#         img_non_hair = (img_non_hair / 255.0) * 2 - 1
+#         ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+#         img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+#         ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+#         pose1 = self.pose[random_number1]
+#         pose1 = torch.tensor(pose1)
+#         pose2 = self.pose[random_number2]
+#         pose2 = torch.tensor(pose2)
+#         hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+#         hair_num = [0, 2, 6, 14, 18, 21]
+#         img_hair_stack = []
+#         for i in hair_num:
+#             img_hair = cv2.imread(hair_path)
+#             img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+#             img_hair = cv2.resize(img_hair, (512, 512))
+#             img_hair = (img_hair / 255.0) * 2 - 1
+#             img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+#             img_hair_stack.append(img_hair)
+#         img_hair = torch.stack(img_hair_stack)
+
+#         return {
+#             'hair_pose': pose1,
+#             'img_hair': img_hair,
+#             'bald_pose': pose2,
+#             'img_non_hair': img_non_hair,
+#             'ref_hair': ref_hair
+#         }
+
+#     def __len__(self):
+#         return self.len
+
+class myDataset(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+    def __init__(self, train_data_dir):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        self.ref_path = os.path.join(train_data_dir, "ref_hair")
+
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)
+        self.pose = np.load(self.pose_path)
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 120)
+        random_number2 = random.randrange(0, 120)
+        while random_number2==random_number1:
+            random_number2 = random.randrange(0, 120)
+        name = self.lists[index]
+
+        hair_path = os.path.join(self.img_path, name, str(random_number1)+'.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2)+'.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        ref_path = os.path.join(ref_folder, files[0])
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+        img_hair = (img_hair/255.0)* 2 - 1
+        img_non_hair = (img_non_hair/255.0)
+        ref_hair = (ref_hair/255.0)* 2 - 1
+
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)  
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)  
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)  
+
+        pose1 = self.pose[random_number1]
+        pose1 = torch.tensor(pose1)   
+        pose2 = self.pose[random_number2]
+        pose2 = torch.tensor(pose2)      
+        
+        return {
+            'hair_pose': pose1, 
+            'img_hair':img_hair, 
+            'bald_pose': pose2, 
+            'img_non_hair':img_non_hair, 
+            'ref_hair':ref_hair
+            }  
+    
+    def __len__(self):
+        return self.len
+    
+
+# class myDataset_unet(data.Dataset):
+#     """Custom data.Dataset compatible with data.DataLoader."""
+
+# class myDataset_unet(data.Dataset):
+#     """Custom data.Dataset compatible with data.DataLoader."""
+
+#     def __init__(self, train_data_dir, frame_num=6):
+#         self.img_path = os.path.join(train_data_dir, "hair")
+#         # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+#         # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+#         # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+#         self.pose_path = os.path.join(train_data_dir, "pose.npy")
+#         self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+#         self.ref_path = os.path.join(train_data_dir, "reference")
+#         self.lists = os.listdir(self.img_path)
+#         self.len = len(self.lists)-10
+#         self.pose = np.load(self.pose_path)
+#         self.frame_num = frame_num
+#         #self.pose = np.random.randn(12, 4)
+
+#     def __getitem__(self, index):
+#         """Returns one data pair (source and target)."""
+#         # seq_len, fea_dim
+#         random_number1 = random.randrange(0, 21)
+#         random_number2 = random.randrange(0, 21)
+
+#         while random_number2 == random_number1:
+#             random_number2 = random.randrange(0, 21)
+#         name = self.lists[index]
+
+#         random_number1 = random_number1
+#         #* 10
+#         #random_number2 = random_number2 * 10
+
+#         random_number2 = random_number1
+
+        
+#         non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+#         ref_folder = os.path.join(self.ref_path, name)
+#         ref_folder = os.path.join(self.img_path, name)
+
+#         files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+#         #ref_path = os.path.join(ref_folder, files[0])
+#         ref_path = os.path.join(ref_folder, '0.jpg')
+        
+#         img_non_hair = cv2.imread(non_hair_path)
+#         ref_hair = cv2.imread(ref_path)
+
+        
+#         img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+#         ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+#         img_non_hair = cv2.resize(img_non_hair, (512, 512))
+#         ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+#         img_non_hair = (img_non_hair / 255.0) * 2 - 1
+#         ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+#         img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+#         ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+#         pose1 = self.pose[random_number1]
+#         pose1 = torch.tensor(pose1)
+#         pose2 = self.pose[random_number2]
+#         pose2 = torch.tensor(pose2)
+#         hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+#         hair_num = [0, 2, 6, 14, 18, 21]
+#         img_hair_stack = []
+#         # begin = random.randrange(0, 21-self.frame_num)
+#         # hair_num = [i+begin for i in range(self.frame_num)]
+#         for i in hair_num:
+#             img_hair = cv2.imread(hair_path)
+#             img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+#             img_hair = cv2.resize(img_hair, (512, 512))
+#             img_hair = (img_hair / 255.0) * 2 - 1
+#             img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+#             img_hair_stack.append(img_hair)
+#         img_hair = torch.stack(img_hair_stack)
+
+#         return {
+#             'hair_pose': pose1,
+#             'img_hair': img_hair,
+#             'bald_pose': pose2,
+#             'img_non_hair': img_non_hair,
+#             'ref_hair': ref_hair
+#         }
+
+#     def __len__(self):
+#         return self.len
+
+class myDataset_unet(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)
+        self.pose = np.load(self.pose_path)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        # while random_number2 == random_number1:
+        #     random_number2 = random.randrange(0, 21)
+        name = self.lists[index]
+
+        #random_number1 = random_number1
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.img_path, name)
+
+        #files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        # img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        # ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        # pose1 = self.pose[random_number1]
+        # pose1 = torch.tensor(pose1)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        # polars = self.polars_rad[random_number1]
+        # polars = torch.tensor(polars).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        pose = self.pose[random_number1]
+        pose = torch.tensor(pose)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            # 'img_non_hair': img_non_hair,
+            'img_ref': ref_hair,
+            'pose': pose,
+            # 'polars': polars,
+            # 'azimuths': azimuths,
+        }
+
+    def __len__(self):
+        return self.len-10
+
+class myDataset_sv3d(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number3 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        while random_number3 == random_number1:
+            random_number3 = random.randrange(0, 21)
+
+        # while random_number3 == random_number1:
+        #     random_number3 = random.randrange(0, 21)
+        name = self.lists[index]
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        #hair_path2 = os.path.join(self.img_path, name, str(random_number3) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        #non_hair_path2 = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        #non_hair_path3 = os.path.join(self.non_hair_path, name, str(random_number3) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        #img_hair2 = cv2.imread(hair_path2)
+        img_non_hair = cv2.imread(non_hair_path)
+        #img_non_hair2 = cv2.imread(non_hair_path2)
+        #img_non_hair3 = cv2.imread(non_hair_path3)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        #img_non_hair2 = cv2.cvtColor(img_non_hair2, cv2.COLOR_BGR2RGB)
+        #img_non_hair3 = cv2.cvtColor(img_non_hair3, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        #img_non_hair2 = cv2.resize(img_non_hair2, (512, 512))
+        #img_non_hair3 = cv2.resize(img_non_hair3, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        #img_non_hair2 = (img_non_hair2 / 255.0) * 2 - 1
+        #img_non_hair3 = (img_non_hair3 / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        #img_non_hair2 = torch.tensor(img_non_hair2).permute(2, 0, 1)
+        #img_non_hair3 = torch.tensor(img_non_hair3).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        pose2 = self.pose[random_number3] 
+        pose2 = torch.tensor(pose2)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        # img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+
+        #img_hair2 = cv2.cvtColor(img_hair2, cv2.COLOR_BGR2RGB)
+        #img_hair2 = cv2.resize(img_hair2, (512, 512))
+        #img_hair2 = (img_hair2 / 255.0) * 2 - 1
+        #img_hair2 = torch.tensor(img_hair2).permute(2, 0, 1)
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            #'img_hair2': img_hair2,
+            # 'bald_pose': pose2,
+            #'pose': pose,
+            #'pose2': pose2,
+            'img_non_hair': img_non_hair,
+            #'img_non_hair2': img_non_hair2,
+            #'img_non_hair3': img_non_hair3,
+            'ref_hair': ref_hair,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d2(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        # while random_number2 == random_number1:
+        #     random_number2 = random.randrange(0, 21)
+        name = self.lists[index]
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+    
+class myDataset_sv3d_temporal(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+        
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21-10)
+        # random_number3 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        while random_number3 == random_number1:
+            random_number3 = random.randrange(0, 21)
+
+        # while random_number3 == random_number1:
+        #     random_number3 = random.randrange(0, 21)
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        for i in range(10):
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(1))
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(1))
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        hair_path2 = os.path.join(self.img_path, name, str(random_number3) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        non_hair_path2 = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        non_hair_path3 = os.path.join(self.non_hair_path, name, str(random_number3) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        img_hair2 = cv2.imread(hair_path2)
+        img_non_hair = cv2.imread(non_hair_path)
+        img_non_hair2 = cv2.imread(non_hair_path2)
+        img_non_hair3 = cv2.imread(non_hair_path3)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        img_non_hair2 = cv2.cvtColor(img_non_hair2, cv2.COLOR_BGR2RGB)
+        img_non_hair3 = cv2.cvtColor(img_non_hair3, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        img_non_hair2 = cv2.resize(img_non_hair2, (512, 512))
+        img_non_hair3 = cv2.resize(img_non_hair3, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        img_non_hair2 = (img_non_hair2 / 255.0) * 2 - 1
+        img_non_hair3 = (img_non_hair3 / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        img_non_hair2 = torch.tensor(img_non_hair2).permute(2, 0, 1)
+        img_non_hair3 = torch.tensor(img_non_hair3).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        pose2 = self.pose[random_number3] 
+        pose2 = torch.tensor(pose2)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        # img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+
+        img_hair2 = cv2.cvtColor(img_hair2, cv2.COLOR_BGR2RGB)
+        img_hair2 = cv2.resize(img_hair2, (512, 512))
+        img_hair2 = (img_hair2 / 255.0) * 2 - 1
+        img_hair2 = torch.tensor(img_hair2).permute(2, 0, 1)
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            'img_hair2': img_hair2,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'pose2': pose2,
+            'img_non_hair': img_non_hair,
+            'img_non_hair2': img_non_hair2,
+            'img_non_hair3': img_non_hair3,
+            'ref_hair': ref_hair,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+    
+class myDataset_sv3d_simple(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        # self.ref_path = os.path.join(train_data_dir, "reference")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        # while random_number2 == random_number1:
+        #     random_number2 = random.randrange(0, 21)
+        name = self.lists[index]
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        random_number2 = random_number1
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        x = torch.tensor(self.x[random_number1])
+        y = torch.tensor(self.y[random_number1])
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+    
+class myDataset_sv3d_simple_ori(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        # while random_number2 == random_number1:
+        #     random_number2 = random.randrange(0, 21)
+        name = self.lists[index]
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number2) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number2] 
+        pose = torch.tensor(pose)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        hair_path = os.path.join(self.img_path, name, str(random_number2) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        x = torch.tensor(self.x[random_number2])
+        y = torch.tensor(self.y[random_number2])
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d_simple_temporal(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        ref_hair = self.read_img(ref_path)
+        for i in range(12):
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            #hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+    
+class myDataset_sv3d_simple_temporal2(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement'
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "hair_good")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+
+        self.pose_path2 = os.path.join(train_data_dir2, "pose.npy")
+        self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair")
+        self.ref_path2 = os.path.join(train_data_dir2, "reference")
+
+        self.lists = os.listdir(self.img_path2)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+    
+    def read_ref_img(self, path):
+        img = cv2.imread(path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = hair_transform(image=img)['image']
+        # img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def reference_lists(self, reference_num, root):
+        stacks = []
+        invalid = []
+        for i in range(12):
+            if (reference_num-6+i)<0:
+                invalid.append(reference_num-5+i+21)
+            else:
+                invalid.append((reference_num-5+i)%21)
+        for i in range(21):
+            if i in invalid:
+                continue
+            else:
+                stacks.append(os.path.join(root, str(i)+'.jpg'))
+        return stacks
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number = random.uniform(0, 1)
+        if random_number<0.5:
+            non_hair_root = self.non_hair_path2
+            img_path = self.img_path2
+        else:
+            non_hair_root = self.non_hair_path
+            img_path = self.img_path
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index].split('.')[0]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        # non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        # files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3] + self.reference_lists(random_number2, os.path.join(self.img_path, name))[:5]
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        ref_hair = self.read_ref_img(ref_path)
+        for i in range(12):
+            #non_hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg')
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            # hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d_simple_temporal3(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement'
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "non-hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+
+        self.pose_path2 = os.path.join(train_data_dir2, "pose.npy")
+        self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair")
+        self.ref_path2 = os.path.join(train_data_dir2, "reference")
+
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number = random.uniform(0, 1)
+        if random_number<0.5:
+            non_hair_root = self.non_hair_path2
+            img_path = self.img_path2
+        else:
+            non_hair_root = self.non_hair_path
+            img_path = self.img_path
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        # non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        ref_hair = self.read_img(ref_path)
+        for i in range(12):
+            non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg')
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            # hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+
+class myDataset_sv3d_simple_temporal_controlnet_without_pose(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement'
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "non-hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+
+        self.pose_path2 = os.path.join(train_data_dir2, "pose.npy")
+        self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair")
+        self.ref_path2 = os.path.join(train_data_dir2, "reference")
+
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number = random.uniform(0, 1)
+        if random_number<0.5:
+            non_hair_root = self.non_hair_path2
+            img_path = self.img_path2
+        else:
+            non_hair_root = self.non_hair_path
+            img_path = self.img_path
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        ref_hair = self.read_img(ref_path)
+        for i in range(12):
+            non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg')
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+    
+    
+class myDataset_sv3d_simple_temporal_controlnet(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement'
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "non-hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+
+        self.pose_path2 = os.path.join(train_data_dir2, "pose.npy")
+        self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair")
+        self.ref_path2 = os.path.join(train_data_dir2, "reference")
+
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number = random.uniform(0, 1)
+        if random_number<0.5:
+            non_hair_root = self.non_hair_path2
+            img_path = self.img_path2
+        else:
+            non_hair_root = self.non_hair_path
+            img_path = self.img_path
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        ref_hair = self.read_img(ref_path)
+        non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg')
+        img_non_hair = self.read_img(non_hair_path)
+        hair_path = os.path.join(img_path, name, str(random_number1) + '.jpg')
+        img_hair= self.read_img(hair_path)
+        x = self.x[random_number1]
+        y = self.y[random_number1]
+            
+
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d_simple_temporal_pose(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        random_number = random.randint(0, 1)
+        if random_number==0:
+            img_path = self.img_path
+        else:
+            img_path = self.non_hair_path
+        non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        ref_hair = self.read_img(ref_path)
+        
+
+        for i in range(12):
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+
+class myDataset_sv3d_simple_temporal_random_reference(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21-12)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        ref_hair = self.read_img(ref_path)
+        for i in range(12):
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d_simple_random_reference(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "hair_good")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+        # self.lists = os.listdir(self.img_path2)
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+
+        # while random_number2 == random_number1:
+        #     random_number2 = random.randrange(0, 21)
+        name = self.lists[index].split('.')[0]
+
+        #random_number1 = random_number1
+        #* 10
+        #random_number2 = random_number2 * 10
+
+        #random_number2 = random_number1
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        # print('________')
+        # print(files)
+        # print('++++++++')
+        # print(ref_folder)
+        # print("========")
+        # print(name)
+        # print("********")
+        # print(ref_path)
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        # pose2 = self.pose[random_number2]
+        # pose2 = torch.tensor(pose2)
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_num = [0, 2, 6, 14, 18, 21]
+        # img_hair_stack = []
+        # polar = self.polars_rad[random_number1]
+        # polar = torch.tensor(polar).unsqueeze(0)
+        # azimuths = self.azimuths_rad[random_number1]
+        # azimuths = torch.tensor(azimuths).unsqueeze(0)
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        x = torch.tensor(self.x[random_number1])
+        y = torch.tensor(self.y[random_number1])
+        x2 = torch.tensor(self.x[random_number2])
+        y2 = torch.tensor(self.y[random_number2])
+        # begin = random.randrange(0, 21-self.frame_num)
+        # hair_num = [i+begin for i in range(self.frame_num)]
+        # for i in hair_num:
+        #     img_hair = cv2.imread(hair_path)
+        #     img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        #     img_hair = cv2.resize(img_hair, (512, 512))
+        #     img_hair = (img_hair / 255.0) * 2 - 1
+        #     img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        #     img_hair_stack.append(img_hair)
+        # img_hair = torch.stack(img_hair_stack)
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+            # 'x2': x2,
+            # 'y2': y2,
+            # 'polar': polar,
+            # 'azimuths':azimuths,
+        }
+
+    def __len__(self):
+        return self.len
+
+class myDataset_sv3d_simple_random_reference_controlnet(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        self.img_path2 = os.path.join(train_data_dir, "hair_good")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "multi_reference2")
+        # self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+
+    def read_ref_img(self, path):
+        img = cv2.imread(path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = hair_transform(image=img)
+        # img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+        name = self.lists[index].split('.')[0]
+        
+        # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3]
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+        
+        ref_hair = hair_transform(image=ref_hair)['image']
+        # print(type(ref_hair))
+        # print(ref_hair.keys())
+        # ref_hair = self.read_ref_img(ref_path)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        #hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        # hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg')
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        x = torch.tensor(self.x[random_number1])
+        y = torch.tensor(self.y[random_number1])
+        x2 = torch.tensor(self.x[random_number2])
+        y2 = torch.tensor(self.y[random_number2])
+
+        return {
+            # 'hair_pose': pose1,
+            'img_hair': img_hair,
+            # 'bald_pose': pose2,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+        }
+
+    def __len__(self):
+        return self.len
+
+    
+class myDataset_sv3d_simple_random_reference_stable_hair(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21)
+        random_number2 = random.randrange(0, 21)
+        random_number1 = random_number2
+
+        name = self.lists[index]
+
+
+        
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        ref_path = os.path.join(ref_folder,random.choice(files))
+        img_hair = cv2.imread(hair_path)
+        img_non_hair = cv2.imread(non_hair_path)
+        ref_hair = cv2.imread(ref_path)
+
+        
+        img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB)
+        ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB)
+
+        
+        img_non_hair = cv2.resize(img_non_hair, (512, 512))
+        ref_hair = cv2.resize(ref_hair, (512, 512))
+
+        
+        img_non_hair = (img_non_hair / 255.0) * 2 - 1
+        ref_hair = (ref_hair / 255.0) * 2 - 1
+
+        
+        img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1)
+        ref_hair = torch.tensor(ref_hair).permute(2, 0, 1)
+
+        pose = self.pose[random_number1] 
+        pose = torch.tensor(pose)
+        hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg')
+        img_hair = cv2.imread(hair_path)
+        img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB)
+        img_hair = cv2.resize(img_hair, (512, 512))
+        img_hair = (img_hair / 255.0) * 2 - 1
+        img_hair = torch.tensor(img_hair).permute(2, 0, 1)
+        x = torch.tensor(self.x[random_number1])
+        y = torch.tensor(self.y[random_number1])
+
+        return {
+            'img_hair': img_hair,
+            'pose': pose,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+        }
+
+    def __len__(self):
+        return self.len
+    
+
+class myDataset_sv3d_simple_temporal_small_squence(data.Dataset):
+    """Custom data.Dataset compatible with data.DataLoader."""
+
+    def __init__(self, train_data_dir, frame_num=6):
+        self.img_path = os.path.join(train_data_dir, "hair")
+        # self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        # self.non_hair_path = os.path.join(train_data_dir, "no_hair")
+        # self.ref_path = os.path.join(train_data_dir, "ref_hair")
+        self.pose_path = os.path.join(train_data_dir, "pose.npy")
+        self.non_hair_path = os.path.join(train_data_dir, "non-hair")
+        self.ref_path = os.path.join(train_data_dir, "reference")
+        self.lists = os.listdir(self.img_path)
+        self.len = len(self.lists)-10
+        self.pose = np.load(self.pose_path)
+        self.frame_num = frame_num
+        #self.pose = np.random.randn(12, 4)
+        elevations_deg = [-0.05/2*np.pi*360] * 21
+        azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360
+        Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)]
+        for i in Face_yaws:
+            if i<0:
+                i = 2*np.pi+i
+            i = i/2*np.pi*360
+        face_yaws = [Face_yaws[0]]
+        for i in range(20):
+            face_yaws.append(Face_yaws[3*i+2])
+        self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg]
+        self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg]
+        self.azimuths_rad[:-1].sort()
+        x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)]
+        y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)]
+        self.x = [x[0]]
+        self.y = [y[0]]
+        for i in range(20):
+            self.x.append(x[i*3+2])
+            self.y.append(y[i*3+2])
+            
+    def read_img(self, path):
+        img = cv2.imread(path)
+
+        img = cv2.resize(img, (512, 512))
+        img = (img / 255.0) * 2 - 1
+        img = torch.tensor(img).permute(2, 0, 1)
+        return img
+
+    def __getitem__(self, index):
+        """Returns one data pair (source and target)."""
+        # seq_len, fea_dim
+        random_number1 = random.randrange(0, 21-6)
+        random_number2 = random.randrange(0, 21)
+
+        name = self.lists[index]
+        x_stack = []
+        y_stack = []
+        img_non_hair_stack = []
+        img_hair_stack = []
+        non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg')
+        ref_folder = os.path.join(self.ref_path, name)
+
+        files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')]
+        # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg')
+        ref_path = os.path.join(ref_folder,files[0])
+        ref_hair = self.read_img(ref_path)
+        for i in range(6):
+            img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0))
+            hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg')
+            img_hair_stack.append(self.read_img(hair_path).unsqueeze(0))
+            x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0))
+            y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0))
+            
+        img_non_hair = torch.cat(img_non_hair_stack, axis=0)
+        img_hair = torch.cat(img_hair_stack, axis=0)
+        x = torch.cat(x_stack, axis=0)
+        y = torch.cat(y_stack, axis=0)
+
+        return {
+            'img_hair': img_hair,
+            'img_non_hair': img_non_hair,
+            'ref_hair': ref_hair,
+            'x': x,
+            'y': y,
+
+        }
+
+    def __len__(self):
+        return self.len
+
+
+if __name__ == "__main__":
+
+    train_dataset = myDataset("./data")
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=1,
+        num_workers=1,
+    )
+
+    for epoch in range(0, len(train_dataset) + 1):
+        for step, batch in enumerate(train_dataloader):
+            print("batch[hair_pose]:", batch["hair_pose"])
+            print("batch[img_hair]:", batch["img_hair"])
+            print("batch[bald_pose]:", batch["bald_pose"])
+            print("batch[img_non_hair]:", batch["img_non_hair"])
+            print("batch[ref_hair]:", batch["ref_hair"])
+
+
+
+

download.py

@@ -0,0 +1,8 @@
+from huggingface_hub import snapshot_download
+snapshot_download(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    local_dir="stable-diffusion-v1-5/stable-diffusion-v1-5",
+    local_dir_use_symlinks=False,
+    resume_download=True,
+    allow_patterns=["unet/*","vae/*","tokenizer/*","scheduler/*","model_index.json"]
+)

gradio_app.py

@@ -0,0 +1,379 @@
+import os
+os.environ.setdefault("GRADIO_TEMP_DIR", "/data2/lzliu/tmp/gradio")
+os.environ.setdefault("TMPDIR", "/data2/lzliu/tmp")
+os.makedirs("/data2/lzliu/tmp/gradio", exist_ok=True)
+os.makedirs("/data2/lzliu/tmp", exist_ok=True)
+
+
+# 其余保持不变
+
+
+import logging
+import gradio as gr
+import torch
+import os
+import uuid
+from test_stablehairv2 import log_validation
+from test_stablehairv2 import UNet3DConditionModel, ControlNetModel, CCProjection
+from test_stablehairv2 import AutoTokenizer, CLIPVisionModelWithProjection, AutoencoderKL, UNet2DConditionModel
+from omegaconf import OmegaConf
+import numpy as np
+import cv2
+from test_stablehairv2 import _maybe_align_image
+from HairMapper.hair_mapper_run import bald_head
+
+import base64
+
+with open("imgs/background.jpg", "rb") as f:
+    b64_img = base64.b64encode(f.read()).decode()
+
+
+def inference(id_image, hair_image):
+    os.makedirs("gradio_inputs", exist_ok=True)
+    os.makedirs("gradio_outputs", exist_ok=True)
+
+    id_path = "gradio_inputs/id.png"
+    hair_path = "gradio_inputs/hair.png"
+    id_image.save(id_path)
+    hair_image.save(hair_path)
+
+    # ===== 图像对齐 =====
+    aligned_id = _maybe_align_image(id_path, output_size=1024, prefer_cuda=True)
+    aligned_hair = _maybe_align_image(hair_path, output_size=1024, prefer_cuda=True)
+
+    # 保存对齐结果（方便 Gradio 输出）
+    aligned_id_path = "gradio_outputs/aligned_id.png"
+    aligned_hair_path = "gradio_outputs/aligned_hair.png"
+    cv2.imwrite(aligned_id_path, cv2.cvtColor(aligned_id, cv2.COLOR_RGB2BGR))
+    cv2.imwrite(aligned_hair_path, cv2.cvtColor(aligned_hair, cv2.COLOR_RGB2BGR))
+
+    # ===== 调用 HairMapper 秃头化 =====
+    bald_id_path = "gradio_outputs/bald_id.png"
+    cv2.imwrite(bald_id_path, cv2.cvtColor(aligned_id, cv2.COLOR_RGB2BGR))
+    bald_head(bald_id_path, bald_id_path)
+
+    # ===== 原本的 Args =====
+    class Args:
+        pretrained_model_name_or_path = "./stable-diffusion-v1-5/stable-diffusion-v1-5"
+        model_path = "./trained_model"
+        image_encoder = "openai/clip-vit-large-patch14"
+        controlnet_model_name_or_path = None
+        revision = None
+        output_dir = "gradio_outputs"
+        seed = 42
+        num_validation_images = 1
+        validation_ids = [aligned_id_path]  # 用对齐后的图像
+        validation_hairs = [aligned_hair_path]  # 用对齐后的图像
+        use_fp16 = False
+
+    args = Args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # 初始化 logger
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger = logging.getLogger(__name__)
+
+    # ===== 模型加载（和 main() 对齐） =====
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer",
+                                              revision=args.revision)
+    image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder, revision=args.revision).to(device)
+    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).to(
+        device, dtype=torch.float32)
+
+    infer_config = OmegaConf.load('./configs/inference/inference_v2.yaml')
+
+    unet2 = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, torch_dtype=torch.float32
+    ).to(device)
+    conv_in_8 = torch.nn.Conv2d(8, unet2.conv_in.out_channels, kernel_size=unet2.conv_in.kernel_size,
+                                padding=unet2.conv_in.padding)
+    conv_in_8.requires_grad_(False)
+    unet2.conv_in.requires_grad_(False)
+    torch.nn.init.zeros_(conv_in_8.weight)
+    conv_in_8.weight[:, :4, :, :].copy_(unet2.conv_in.weight)
+    conv_in_8.bias.copy_(unet2.conv_in.bias)
+    unet2.conv_in = conv_in_8
+
+    controlnet = ControlNetModel.from_unet(unet2).to(device)
+    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model.bin"), map_location="cpu")
+    controlnet.load_state_dict(state_dict2, strict=False)
+
+    prefix = "motion_module"
+    ckpt_num = "4140000"
+    save_path = os.path.join(args.model_path, f"{prefix}-{ckpt_num}.pth")
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        args.pretrained_model_name_or_path,
+        save_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(device)
+
+    cc_projection = CCProjection().to(device)
+    state_dict3 = torch.load(os.path.join(args.model_path, "pytorch_model_1.bin"), map_location="cpu")
+    cc_projection.load_state_dict(state_dict3, strict=False)
+
+    from ref_encoder.reference_unet import ref_unet
+    Hair_Encoder = ref_unet.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, low_cpu_mem_usage=False,
+        device_map=None, ignore_mismatched_sizes=True
+    ).to(device)
+    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location="cpu")
+    Hair_Encoder.load_state_dict(state_dict2, strict=False)
+
+    # 推理
+    log_validation(
+        vae, tokenizer, image_encoder, denoising_unet,
+        args, device, logger,
+        cc_projection, controlnet, Hair_Encoder
+    )
+
+    output_video = os.path.join(args.output_dir, "validation", "generated_video_0.mp4")
+
+    # 提取视频帧用于可拖动预览
+    frames_dir = os.path.join(args.output_dir, "frames", uuid.uuid4().hex)
+    os.makedirs(frames_dir, exist_ok=True)
+    cap = cv2.VideoCapture(output_video)
+    frames_list = []
+    idx = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        fp = os.path.join(frames_dir, f"{idx:03d}.png")
+        cv2.imwrite(fp, frame)
+        frames_list.append(fp)
+        idx += 1
+    cap.release()
+
+    max_frames = len(frames_list) if frames_list else 1
+    first_frame = frames_list[0] if frames_list else None
+
+    return aligned_id_path, aligned_hair_path, bald_id_path, output_video, frames_list, gr.update(minimum=1,
+                                                                                                  maximum=max_frames,
+                                                                                                  value=1,
+                                                                                                  step=1), first_frame
+
+
+# Gradio 前端
+# 原 Interface 版本（保留以便回退）
+# demo = gr.Interface(
+#     fn=inference,
+#     inputs=[
+#         gr.Image(type="pil", label="上传身份图（ID Image）"),
+#         gr.Image(type="pil", label="上传发型图（Hair Reference Image）")
+#     ],
+#     outputs=[
+#         gr.Image(type="filepath", label="对齐后的身份图"),
+#         gr.Image(type="filepath", label="对齐后的发型图"),
+#         gr.Image(type="filepath", label="秃头化后的身份图"),
+#         gr.Video(label="生成的视频")
+#     ],
+#     title="StableHairV2 多视角发型迁移",
+#     description="上传身份图和发型参考图，查看对齐结果并生成多视角视频"
+# )
+# if __name__ == "__main__":
+#     demo.launch(server_name="0.0.0.0", server_port=7860)
+
+# Blocks 美化版
+css = f"""
+    html, body {{
+        height: 100%;
+        margin: 0;
+        padding: 0;
+    }}
+    .gradio-container {{
+        width: 100% !important;
+        height: 100% !important;
+        margin: 0 !important;
+        padding: 0 !important;
+        background-image: url("data:image/jpeg;base64,{b64_img}");
+        background-size: cover;
+        background-position: center;
+        background-attachment: fixed; /* 背景固定 */
+    }}
+    #title-card {{
+        background: rgba(255, 255, 255, 0.8);
+        border-radius: 12px;
+        padding: 16px 24px;
+        box-shadow: 0 2px 8px rgba(0,0,0,0.15);
+        margin-bottom: 20px;
+    }}
+    #title-card h2 {{
+        text-align: center;
+        margin: 4px 0 12px 0;
+        font-size: 28px;
+    }}
+    #title-card p {{
+        text-align: center;
+        font-size: 16px;
+        color: #374151;
+    }}
+    .out-card {{
+        border:1px solid #e5e7eb; border-radius:10px; padding:10px;
+        background: rgba(255,255,255,0.85);
+    }}
+    .two-col {{
+        display:grid !important; grid-template-columns: 360px minmax(680px, 1fr); gap:16px
+    }}
+    .left-pane {{min-width: 360px}}
+    .right-pane {{min-width: 680px}}
+    /* Tabs 美化 */
+    .tabs {{
+        background: rgba(255,255,255,0.88);
+        border-radius: 12px;
+        box-shadow: 0 8px 24px rgba(0,0,0,0.08);
+        padding: 8px;
+        border: 1px solid #e5e7eb;
+    }}
+    .tab-nav {{
+        display: flex; gap: 8px; margin-bottom: 8px;
+        background: transparent;
+        border-bottom: 1px solid #e5e7eb;
+        padding-bottom: 6px;
+    }}
+    .tab-nav button {{
+        background: rgba(255,255,255,0.7);
+        border: 1px solid #e5e7eb;
+        backdrop-filter: blur(6px);
+        border-radius: 8px;
+        padding: 6px 12px;
+        color: #111827;
+        transition: all .2s ease;
+    }}
+    .tab-nav button:hover {{
+        transform: translateY(-1px);
+        box-shadow: 0 4px 10px rgba(0,0,0,0.06);
+    }}
+    .tab-nav button[aria-selected="true"] {{
+        background: #4f46e5;
+        color: #fff;
+        border-color: #4f46e5;
+        box-shadow: 0 6px 14px rgba(79,70,229,0.25);
+    }}
+    .tabitem {{
+        background: rgba(255,255,255,0.88);
+        border-radius: 10px;
+        padding: 8px;
+    }}
+    /* 发型库滚动限制容器：固定260px高度，内部可滚动 */
+    #hair_gallery_wrap {{
+        height: 260px !important;
+        overflow-y: scroll !important;
+        overflow-x: auto !important;
+    }}
+    #hair_gallery_wrap .grid, #hair_gallery_wrap .wrap {{
+        height: 100% !important;
+        overflow-y: scroll !important;
+    }}
+    /* 确保画廊本体占满容���高度，避免滚动条落到页面底部 */
+    #hair_gallery {{
+        height: 100% !important;
+    }}
+"""
+
+with gr.Blocks(
+        theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
+        css=css
+) as demo:
+    # ==== 顶部 Panel ====
+    with gr.Group(elem_id="title-card"):
+        gr.Markdown("""
+        <h2 id='title'>StableHairV2 多视角发型迁移</h2>
+        <p>上传身份图与发型参考图，系统将自动完成 <b>对齐 → 秃头化 → 视频生成</b>。</p>
+        """)
+
+    with gr.Row(elem_classes=["two-col"]):
+        with gr.Column(scale=5, min_width=260, elem_classes=["left-pane"]):
+            id_input = gr.Image(type="pil", label="身份图", height=200)
+            hair_input = gr.Image(type="pil", label="发型参考图", height=200)
+
+            with gr.Row():
+                run_btn = gr.Button("开始生成", variant="primary")
+                clear_btn = gr.Button("清空")
+
+            # ========= 发型库（点击即填充到“发型参考图”） =========
+            def _list_imgs(dir_path: str):
+                exts = (".png", ".jpg", ".jpeg", ".webp")
+                # exts = (".jpg")
+                try:
+                    files = [os.path.join(dir_path, f) for f in sorted(os.listdir(dir_path))
+                             if f.lower().endswith(exts)]
+                    return files
+                except Exception:
+                    return []
+
+            hair_list = _list_imgs("hair_resposity")
+
+            with gr.Accordion("发型库（点击选择后自动填充）", open=True):
+                with gr.Group(elem_id="hair_gallery_wrap"):
+                    gallery = gr.Gallery(
+                        value=hair_list,
+                        columns=4, rows=2, allow_preview=True, label="发型库",
+                        elem_id="hair_gallery"
+                    )
+
+            def _pick_hair(evt: gr.SelectData):  # type: ignore[name-defined]
+                i = evt.index if hasattr(evt, 'index') else 0
+                i = 0 if i is None else int(i)
+                if 0 <= i < len(hair_list):
+                    return gr.update(value=hair_list[i])
+                return gr.update()
+
+            gallery.select(_pick_hair, inputs=None, outputs=hair_input)
+
+        with gr.Column(scale=7, min_width=520, elem_classes=["right-pane"]):
+            with gr.Tabs():
+                with gr.TabItem("生成视频"):
+                    with gr.Group(elem_classes=["out-card"]):
+                        video_out = gr.Video(label="生成的视频", height=340)
+                        with gr.Row():
+                            frame_slider = gr.Slider(1, 21, value=1, step=1, label="多视角预览（拖动查看帧）")
+                        frame_preview = gr.Image(type="filepath", label="预览帧", height=260)
+                        frames_state = gr.State([])
+
+                with gr.TabItem("归一化对齐结果"):
+                    with gr.Group(elem_classes=["out-card"]):
+                        with gr.Row():
+                            aligned_id_out = gr.Image(type="filepath", label="对齐后的身份图", height=240)
+                            aligned_hair_out = gr.Image(type="filepath", label="对齐后的发型图", height=240)
+
+                with gr.TabItem("秃头化结果"):
+                    with gr.Group(elem_classes=["out-card"]):
+                        bald_id_out = gr.Image(type="filepath", label="秃头化后的身份图", height=260)
+
+    # 逻辑保持不变
+    run_btn.click(fn=inference,
+                  inputs=[id_input, hair_input],
+                  outputs=[aligned_id_out, aligned_hair_out, bald_id_out,
+                           video_out, frames_state, frame_slider, frame_preview])
+
+
+    def _on_slide(frames, idx):
+        if not frames:
+            return gr.update()
+        i = int(idx) - 1
+        i = max(0, min(i, len(frames) - 1))
+        return gr.update(value=frames[i])
+
+
+    frame_slider.change(_on_slide, inputs=[frames_state, frame_slider], outputs=frame_preview)
+
+
+    def _clear():
+        return None, None, None, None, None
+
+
+    clear_btn.click(_clear, None,
+                    [id_input, hair_input, aligned_id_out, aligned_hair_out, bald_id_out])
+
+if __name__ == "__main__":
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+
+
+
+

requirements.txt

@@ -0,0 +1,67 @@
+accelerate
+albucore==0.0.24
+albumentations==2.0.8
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+av==14.4.0
+Brotli
+certifi
+charset-normalizer
+click==8.2.1
+colorama
+einops==0.8.1
+filelock
+fsspec
+gitdb==4.0.12
+GitPython==3.1.44
+gmpy2
+graphviz==0.20.3
+hf-xet==1.1.5
+huggingface-hub==0.30.0
+idna
+Jinja2
+kornia
+MarkupSafe
+mkl-service==2.4.0
+mkl_fft
+mkl_random
+mpmath
+networkx
+numpy
+omegaconf==2.3.0
+opencv-python==4.11.0.86
+opencv-python-headless==4.11.0.86
+packaging
+peft==0.15.2
+pillow
+platformdirs==4.3.8
+prodigyopt==1.1.2
+protobuf==6.31.1
+psutil
+pydantic==2.11.7
+pydantic_core==2.33.2
+PySocks
+PyYAML
+regex==2024.11.6
+requests
+safetensors
+scipy==1.15.3
+sentencepiece==0.2.0
+sentry-sdk==2.32.0
+setproctitle==1.3.6
+simsimd==6.4.9
+smmap==5.0.2
+stringzilla==3.12.5
+sympy==1.13.1
+tokenizers==0.21.1
+torch==2.5.0
+torchaudio==2.5.0
+torchvision==0.20.0
+torchviz==0.0.3
+tqdm
+transformers==4.52.3
+triton==3.1.0
+typing-inspection==0.4.1
+typing_extensions
+urllib3
+wandb==0.20.1

test_stablehairv2.py

@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+import argparse
+import logging
+import sys
+import os
+import random
+import numpy as np
+import cv2
+import torch
+from PIL import Image
+from transformers import AutoTokenizer, CLIPVisionModelWithProjection
+from diffusers import AutoencoderKL, UniPCMultistepScheduler, UNet2DConditionModel
+from src.models.unet_3d import UNet3DConditionModel
+from ref_encoder.reference_unet import CCProjection
+from ref_encoder.latent_controlnet import ControlNetModel
+from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline as Hair3dPipeline
+from src.utils.util import save_videos_grid
+from omegaconf import OmegaConf
+from HairMapper.hair_mapper_run import bald_head
+
+
+# face align
+def _maybe_align_image(image_path: str, output_size: int, prefer_cuda: bool = True):
+    """Align and crop a face image to FFHQ-style using FFHQFaceAlignment if available.
+    Falls back to simple resize if alignment fails.
+    Returns an RGB uint8 numpy array of shape (H, W, 3).
+    """
+    try:
+        ffhq_dir = os.path.join(os.path.dirname(__file__), 'FFHQFaceAlignment')
+        if ffhq_dir not in sys.path:
+            sys.path.insert(0, ffhq_dir)
+        # Lazy imports to avoid hard dependency if user doesn't enable alignment
+        from lib.landmarks_pytorch import LandmarksEstimation
+        from align import align_crop_image
+
+        # Read image as RGB uint8
+        img_bgr = cv2.imread(image_path, cv2.IMREAD_COLOR)
+        if img_bgr is None:
+            raise RuntimeError(f"Failed to read image: {image_path}")
+        img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype('uint8')
+
+        device = torch.device('cuda' if prefer_cuda and torch.cuda.is_available() else 'cpu')
+        le = LandmarksEstimation(type='2D')
+
+        img_tensor = torch.tensor(np.transpose(img, (2, 0, 1))).float().to(device)
+        with torch.no_grad():
+            landmarks, _ = le.detect_landmarks(img_tensor.unsqueeze(0), detected_faces=None)
+        if len(landmarks) > 0:
+            lm = np.asarray(landmarks[0].detach().cpu().numpy())
+            aligned = align_crop_image(image=img, landmarks=lm, transform_size=output_size)
+            if aligned is None or aligned.size == 0:
+                return cv2.resize(img, (output_size, output_size))
+            return aligned
+        else:
+            return cv2.resize(img, (output_size, output_size))
+    except Exception:
+        # Silent fallback to simple resize on any failure
+        img_bgr = cv2.imread(image_path, cv2.IMREAD_COLOR)
+        img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype('uint8') if img_bgr is not None else None
+        if img is None:
+            raise
+        return cv2.resize(img, (output_size, output_size))
+
+
+def log_validation(
+        vae, tokenizer, image_encoder, denoising_unet,
+        args, device, logger, cc_projection,
+        controlnet, hair_encoder, feature_extractor=None
+):
+    """
+    Run inference on validation pairs and save generated videos.
+    """
+    logger.info("Starting validation inference...")
+
+    # Initialize inference pipeline
+    pipeline = Hair3dPipeline.from_pretrained(
+        args.pretrained_model_name_or_path,
+        image_encoder=image_encoder,
+        feature_extractor=feature_extractor,
+        controlnet=controlnet,
+        vae=vae,
+        tokenizer=tokenizer,
+        denoising_unet=denoising_unet,
+        safety_checker=None,
+        revision=args.revision,
+        torch_dtype=torch.float16 if args.use_fp16 else torch.float32,
+    ).to(device)
+    pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
+    pipeline.set_progress_bar_config(disable=True)
+
+    # Create output directory
+    output_dir = os.path.join(args.output_dir, "validation")
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(output_dir)
+
+    # Generate camera trajectory
+    x_coords = [0.4 * np.sin(2 * np.pi * i / 120) for i in range(60)]
+    y_coords = [-0.05 + 0.3 * np.cos(2 * np.pi * i / 120) for i in range(60)]
+    X = [x_coords[0]]
+    Y = [y_coords[0]]
+    for i in range(20):
+        X.append(x_coords[i * 3 + 2])
+        Y.append(y_coords[i * 3 + 2])
+    x_tensor = torch.tensor(X, dtype=torch.float32).unsqueeze(1).to(device)
+    y_tensor = torch.tensor(Y, dtype=torch.float32).unsqueeze(1).to(device)
+
+    # # Load reference images
+    # id_image = cv2.cvtColor(cv2.imread(args.validation_ids[0]), cv2.COLOR_BGR2RGB)
+    # id_image = cv2.resize(id_image, (512, 512))
+    # Load reference images (optionally align)
+    align_enabled = getattr(args, 'align_before_infer', True)
+    align_size = getattr(args, 'align_size', 1024)
+    prefer_cuda = True if device.type == 'cuda' else False
+    if align_enabled:
+        id_image = _maybe_align_image(args.validation_ids[0], output_size=align_size, prefer_cuda=prefer_cuda)
+    else:
+        id_image = cv2.cvtColor(cv2.imread(args.validation_ids[0]), cv2.COLOR_BGR2RGB)
+        id_image = cv2.resize(id_image, (512, 512))
+
+    # ===== ���� HairMapper ͺͷ�� =====
+    temp_bald_path = os.path.join(args.output_dir, "bald_id.png")
+    cv2.imwrite(temp_bald_path, cv2.cvtColor(id_image, cv2.COLOR_RGB2BGR))  # �������ͼ
+    bald_head(temp_bald_path, temp_bald_path)  # ͺͷ�������DZ���
+    # ���¼���ͺͷͼ�� (RGB)
+    id_image = cv2.cvtColor(cv2.imread(temp_bald_path), cv2.COLOR_BGR2RGB)
+    id_image = cv2.resize(id_image, (512, 512))
+
+    id_list = [id_image for _ in range(12)]
+    if align_enabled:
+        hair_image = _maybe_align_image(args.validation_hairs[0], output_size=align_size, prefer_cuda=prefer_cuda)
+        prompt_img = _maybe_align_image(args.validation_ids[0], output_size=align_size, prefer_cuda=prefer_cuda)
+    else:
+        hair_image = cv2.cvtColor(cv2.imread(args.validation_hairs[0]), cv2.COLOR_BGR2RGB)
+        hair_image = cv2.resize(hair_image, (512, 512))
+        prompt_img = cv2.cvtColor(cv2.imread(args.validation_ids[0]), cv2.COLOR_BGR2RGB)
+        prompt_img = cv2.resize(prompt_img, (512, 512))
+    hair_image = cv2.resize(hair_image, (512, 512))
+    prompt_img = cv2.resize(prompt_img, (512, 512))
+
+    prompt_img = [prompt_img]
+
+    # Perform inference and save videos
+    for idx in range(args.num_validation_images):
+        result = pipeline(
+            prompt="",
+            negative_prompt="",
+            num_inference_steps=30,
+            guidance_scale=1.5,
+            width=512,
+            height=512,
+            controlnet_condition=id_list,
+            controlnet_conditioning_scale=1.0,
+            generator=torch.Generator(device).manual_seed(args.seed),
+            ref_image=hair_image,
+            prompt_img=prompt_img,
+            reference_encoder=hair_encoder,
+            poses=None,
+            x=x_tensor,
+            y=y_tensor,
+            video_length=21,
+            context_frames=12,
+        )
+        video = torch.cat([result.videos, result.videos], dim=0)
+        video_path = os.path.join(output_dir, f"generated_video_{idx}.mp4")
+        save_videos_grid(video, video_path, n_rows=5, fps=24)
+        logger.info(f"Saved generated video: {video_path}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Inference script for 3D hairstyle generation"
+    )
+    parser.add_argument(
+        "--pretrained_model_name_or_path", type=str, required=True,
+        help="Path or ID of the pretrained pipeline"
+    )
+    parser.add_argument(
+        "--model_path", type=str, required=True,
+        help="Path or ID of the pretrained pipeline"
+    )
+    parser.add_argument(
+        "--image_encoder", type=str, required=True,
+        help="Path or ID of the CLIP vision encoder"
+    )
+    parser.add_argument(
+        "--controlnet_model_name_or_path", type=str, default=None,
+        help="Path or ID of the ControlNet model"
+    )
+    parser.add_argument(
+        "--revision", type=str, default=None,
+        help="Model revision or Git reference"
+    )
+    parser.add_argument(
+        "--output_dir", type=str, default="inference_output",
+        help="Directory to save inference results"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42,
+        help="Random seed for reproducibility"
+    )
+    parser.add_argument(
+        "--num_validation_images", type=int, default=3,
+        help="Number of videos to generate per input pair"
+    )
+    parser.add_argument(
+        "--validation_ids", type=str, nargs='+', required=True,
+        help="Path(s) to identity conditioning images"
+    )
+    parser.add_argument(
+        "--validation_hairs", type=str, nargs='+', required=True,
+        help="Path(s) to hairstyle reference images"
+    )
+    parser.add_argument(
+        "--use_fp16", action="store_true",
+        help="Enable fp16 inference"
+    )
+    parser.add_argument(
+        "--align_before_infer", action="store_true", default=True,
+        help="Align and crop input images to FFHQ style before inference"
+    )
+    parser.add_argument(
+        "--align_size", type=int, default=1024,
+        help="Output size for aligned images when alignment is enabled"
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    # Setup device and logger
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger = logging.getLogger(__name__)
+
+    # Set random seed
+    torch.manual_seed(args.seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(args.seed)
+
+    # Load models
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="tokenizer",
+        revision=args.revision
+    )
+    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+        args.image_encoder,
+        revision=args.revision
+    ).to(device)
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="vae",
+        revision=args.revision
+    ).to(device)
+
+    infer_config = OmegaConf.load('./configs/inference/inference_v2.yaml')
+
+    unet2 = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", use_safetensors=True, revision=args.revision,
+        torch_dtype=torch.float16
+    ).to(device)
+    conv_in_8 = torch.nn.Conv2d(8, unet2.conv_in.out_channels, kernel_size=unet2.conv_in.kernel_size,
+                                padding=unet2.conv_in.padding)
+    conv_in_8.requires_grad_(False)
+    unet2.conv_in.requires_grad_(False)
+    torch.nn.init.zeros_(conv_in_8.weight)
+    conv_in_8.weight[:, :4, :, :].copy_(unet2.conv_in.weight)
+    conv_in_8.bias.copy_(unet2.conv_in.bias)
+    unet2.conv_in = conv_in_8
+
+    # Load or initialize ControlNet
+    controlnet = ControlNetModel.from_unet(unet2).to(device)
+    # state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model.bin"), map_location=torch.device('cpu'))
+    # state_dict2 = torch.load(args.model_path, map_location=torch.device('cpu'))
+    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model.bin"), map_location=torch.device('cpu'))
+    controlnet.load_state_dict(state_dict2, strict=False)
+
+    # Load 3D UNet motion module
+    prefix = "motion_module"
+    ckpt_num = "4140000"
+    save_path = os.path.join(args.model_path, f"{prefix}-{ckpt_num}.pth")
+
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        args.pretrained_model_name_or_path,
+        save_path,
+        subfolder="unet",
+
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(device)
+
+    # Load projection and hair encoder
+    cc_projection = CCProjection().to(device)
+    state_dict3 = torch.load(os.path.join(args.model_path, "pytorch_model_1.bin"), map_location=torch.device('cpu'))
+    cc_projection.load_state_dict(state_dict3, strict=False)
+
+    from ref_encoder.reference_unet import ref_unet
+    Hair_Encoder = ref_unet.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, low_cpu_mem_usage=False,
+        device_map=None, ignore_mismatched_sizes=True
+    ).to(device)
+
+    state_dict2 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location=torch.device('cpu'))
+    # state_dict2 = torch.load(os.path.join('/home/jichao.zhang/code/3dhair/train_sv3d/checkpoint-30000/', "pytorch_model.bin"))
+    Hair_Encoder.load_state_dict(state_dict2, strict=False)
+
+    # Run validation inference
+    log_validation(
+        vae, tokenizer, image_encoder, denoising_unet,
+        args, device, logger,
+        cc_projection, controlnet, Hair_Encoder
+    )
+
+
+if __name__ == "__main__":
+    main()