init upload

README.md

@@ -11,45 +11,62 @@ from diffusers.utils import export_to_video
 from torchvision import transforms
 import os
 import cv2
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 import numpy as np
 
-pretrained_model_name_or_path = "./wan_t2v"
-transformer_t2v = WanTransformer3DModel.from_pretrained(pretrained_model_name_or_path, subfolder='transformer')
 
-text_encoder = UMT5EncoderModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder',
-                                                torch_dtype=torch.bfloat16)
-
-pipe = WanPipeline.from_pretrained(
-    pretrained_model_name_or_path,
-    transformer=transformer_t2v,
-    text_encoder=text_encoder,
-)
-
-negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+from pathlib import Path
+import json
+from safetensors.torch import safe_open
 
 device = "cuda"
 seed = 0
 
+# TODO: impl AutoencoderKLWan
+vae = vae.from_pretrained("StevenZhang/Wan2.1-VAE_Diff")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+vae = vae.to(device)
+
+# TODO: impl FlowDPMSolverMultistepScheduler
+scheduler = UniPCMultistepScheduler(prediction_type='flow_prediction', use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=1.0)
+
+text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl")
+
+# 14B
+transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-14B-Diff', torch_dtype=torch.bfloat16)
+# transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-1.3B-Diff', torch_dtype=torch.bfloat16)
+
+components = {
+    "transformer": transformer,
+    "vae": vae,
+    "scheduler": scheduler,
+    "text_encoder": text_encoder,
+    "tokenizer": tokenizer,
+}
+pipe = WanPipeline(**components)
+
+pipe.to(device)
+
+negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+
 generator = torch.Generator(device=device).manual_seed(seed)
 inputs = {
     "prompt": "两只拟人化的猫咪身穿舒适的拳击装备，戴着鲜艳的手套，在聚光灯照射的舞台上激烈对战",
-    "negative_prompt": negative_prompt,
+    "negative_prompt": negative_prompt, # TODO
     "generator": generator,
     "num_inference_steps": 50,
-    "flow_shift": 5.0,
+    "flow_shift": 3.0,
     "guidance_scale": 5.0,
-    "height": 720,
-    "width": 1280,
+    "height": 480,
+    "width": 832,
     "num_frames": 81,
     "max_sequence_length": 512,
     "output_type": "np"
 }
 
-pipe.enable_model_cpu_offload()
-
 video = pipe(**inputs).frames[0]
 
-export_to_video(video, "output.mp4", fps=16)
+print(video.shape)
 
-```
+export_to_video(video, "output.mp4", fps=16)
+```

wan_t2v_fp32_example.py

@@ -6,43 +6,61 @@ from diffusers.utils import export_to_video
 from torchvision import transforms
 import os
 import cv2
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 import numpy as np
 
-pretrained_model_name_or_path = "./wan_t2v"
-transformer_t2v = WanTransformer3DModel.from_pretrained(pretrained_model_name_or_path, subfolder='transformer')
 
-text_encoder = UMT5EncoderModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder',
-                                                torch_dtype=torch.bfloat16)
-
-pipe = WanPipeline.from_pretrained(
-    pretrained_model_name_or_path,
-    transformer=transformer_t2v,
-    text_encoder=text_encoder,
-)
-
-negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+from pathlib import Path
+import json
+from safetensors.torch import safe_open
 
 device = "cuda"
 seed = 0
 
+# TODO: impl AutoencoderKLWan
+vae = vae.from_pretrained("StevenZhang/Wan2.1-VAE_Diff")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+vae = vae.to(device)
+
+# TODO: impl FlowDPMSolverMultistepScheduler
+scheduler = UniPCMultistepScheduler(prediction_type='flow_prediction', use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=1.0)
+
+text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl")
+
+# 14B
+transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-14B-Diff', torch_dtype=torch.bfloat16)
+# transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-1.3B-Diff', torch_dtype=torch.bfloat16)
+
+components = {
+    "transformer": transformer,
+    "vae": vae,
+    "scheduler": scheduler,
+    "text_encoder": text_encoder,
+    "tokenizer": tokenizer,
+}
+pipe = WanPipeline(**components)
+
+pipe.to(device)
+
+negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+
 generator = torch.Generator(device=device).manual_seed(seed)
 inputs = {
     "prompt": "两只拟人化的猫咪身穿舒适的拳击装备，戴着鲜艳的手套，在聚光灯照射的舞台上激烈对战",
-    "negative_prompt": negative_prompt,
+    "negative_prompt": negative_prompt, # TODO
     "generator": generator,
     "num_inference_steps": 50,
-    "flow_shift": 5.0,
+    "flow_shift": 3.0,
     "guidance_scale": 5.0,
-    "height": 720,
-    "width": 1280,
+    "height": 480,
+    "width": 832,
     "num_frames": 81,
     "max_sequence_length": 512,
     "output_type": "np"
 }
 
-pipe.enable_model_cpu_offload()
-
 video = pipe(**inputs).frames[0]
 
-export_to_video(video, "output.mp4", fps=16)
+print(video.shape)
+
+export_to_video(video, "output.mp4", fps=16)