Update test_stablehairv2.py

test_stablehairv2.py

@@ -9,7 +9,7 @@ import cv2
 import torch
 from PIL import Image
 from transformers import AutoTokenizer, CLIPVisionModelWithProjection
-from diffusers import AutoencoderKL, UniPCMultistepScheduler, UNet2DConditionModel
+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
@@ -61,11 +61,10 @@ def _maybe_align_image(image_path: str, output_size: int, prefer_cuda: bool = Tr
             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
+    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.
@@ -94,21 +93,14 @@ def log_validation(
 
     print(output_dir)
 
-    # Speed/length overrides via env/args
-    import os as _os
-    steps = int(_os.getenv('SH_STEPS', getattr(args, 'num_inference_steps', 30)))
-    gscale = float(_os.getenv('SH_GUIDANCE', getattr(args, 'guidance_scale', 1.5)))
-    vlen = int(_os.getenv('SH_VIDEO_LENGTH', getattr(args, 'video_length', 21)))
-    # 统一时序长度：上下文帧数始终等于视频帧数（不再读取 SH_CONTEXT_FRAMES）
-    cframes = int(_os.getenv('SH_CFRAMES', getattr(args, 'cframes', 12)))
-    print("[cfg]推理步数：",steps)
-    print("[cfg]guidance_scale:",gscale)
-    print("[cfg]视频帧数：",vlen)
-    print("[cfg]cframes:",cframes)
-    # Generate camera trajectory with exactly vlen frames
-    angles = np.linspace(0, 2 * np.pi, vlen, endpoint=False)
-    X = 0.4 * np.sin(angles)
-    Y = -0.05 + 0.3 * np.cos(angles)
+    # 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)
 
@@ -132,8 +124,8 @@ def log_validation(
     # ���¼���ͺͷͼ�� (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(cframes)]
+    
+    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)
@@ -144,17 +136,16 @@ def log_validation(
         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=steps,
-            guidance_scale=gscale,
+            num_inference_steps=30,
+            guidance_scale=1.5,
             width=512,
             height=512,
             controlnet_condition=id_list,
@@ -166,8 +157,8 @@ def log_validation(
             poses=None,
             x=x_tensor,
             y=y_tensor,
-            video_length=vlen,
-            context_frames=cframes,
+            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")
@@ -269,15 +260,13 @@ def main():
     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
+        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 = 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.weight[:,:4,:,:].copy_(unet2.conv_in.weight)
     conv_in_8.bias.copy_(unet2.conv_in.bias)
     unet2.conv_in = conv_in_8
 
@@ -308,12 +297,11 @@ def main():
 
     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
+            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"))
+    #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
@@ -323,6 +311,5 @@ def main():
         cc_projection, controlnet, Hair_Encoder
     )
 
-
 if __name__ == "__main__":
     main()