Spaces:
Paused
Paused
Image to video: updated script.
Browse files- requirements.txt +3 -1
- xora/examples/image_to_video.py +127 -44
requirements.txt
CHANGED
@@ -3,4 +3,6 @@ diffusers==0.28.2
|
|
3 |
transformers==4.44.2
|
4 |
sentencepiece>=0.1.96
|
5 |
accelerate
|
6 |
-
einops
|
|
|
|
|
|
3 |
transformers==4.44.2
|
4 |
sentencepiece>=0.1.96
|
5 |
accelerate
|
6 |
+
einops
|
7 |
+
matplotlib
|
8 |
+
opencv-python
|
xora/examples/image_to_video.py
CHANGED
@@ -9,6 +9,11 @@ from transformers import T5EncoderModel, T5Tokenizer
|
|
9 |
import safetensors.torch
|
10 |
import json
|
11 |
import argparse
|
|
|
|
|
|
|
|
|
|
|
12 |
|
13 |
def load_vae(vae_dir):
|
14 |
vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
|
@@ -34,78 +39,156 @@ def load_scheduler(scheduler_dir):
|
|
34 |
scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
|
35 |
return RectifiedFlowScheduler.from_config(scheduler_config)
|
36 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
37 |
def main():
|
38 |
-
|
39 |
-
|
40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
41 |
args = parser.parse_args()
|
42 |
|
43 |
# Paths for the separate mode directories
|
44 |
-
|
45 |
-
unet_dir =
|
46 |
-
vae_dir =
|
47 |
-
scheduler_dir =
|
48 |
|
49 |
# Load models
|
50 |
vae = load_vae(vae_dir)
|
51 |
unet = load_unet(unet_dir)
|
52 |
scheduler = load_scheduler(scheduler_dir)
|
53 |
-
|
54 |
-
# Patchifier (remains the same)
|
55 |
patchifier = SymmetricPatchifier(patch_size=1)
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
|
60 |
# Use submodels for the pipeline
|
61 |
submodel_dict = {
|
62 |
-
"transformer": unet,
|
63 |
"patchifier": patchifier,
|
64 |
-
"text_encoder":
|
65 |
-
"tokenizer":
|
66 |
"scheduler": scheduler,
|
67 |
"vae": vae,
|
68 |
}
|
69 |
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
|
89 |
-
|
90 |
|
91 |
-
#
|
92 |
images = pipeline(
|
93 |
-
num_inference_steps=num_inference_steps,
|
94 |
-
num_images_per_prompt=num_images_per_prompt,
|
95 |
-
guidance_scale=guidance_scale,
|
96 |
-
generator=
|
97 |
output_type="pt",
|
98 |
callback_on_step_end=None,
|
99 |
-
height=height,
|
100 |
-
width=width,
|
101 |
-
num_frames=num_frames,
|
102 |
-
frame_rate=frame_rate,
|
103 |
**sample,
|
104 |
is_video=True,
|
105 |
vae_per_channel_normalize=True,
|
|
|
106 |
).images
|
107 |
|
108 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
109 |
|
110 |
if __name__ == "__main__":
|
111 |
main()
|
|
|
9 |
import safetensors.torch
|
10 |
import json
|
11 |
import argparse
|
12 |
+
from xora.utils.conditioning_method import ConditioningMethod
|
13 |
+
import os
|
14 |
+
import numpy as np
|
15 |
+
import cv2
|
16 |
+
from PIL import Image
|
17 |
|
18 |
def load_vae(vae_dir):
|
19 |
vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
|
|
|
39 |
scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
|
40 |
return RectifiedFlowScheduler.from_config(scheduler_config)
|
41 |
|
42 |
+
def center_crop_and_resize(frame, target_height, target_width):
|
43 |
+
h, w, _ = frame.shape
|
44 |
+
aspect_ratio_target = target_width / target_height
|
45 |
+
aspect_ratio_frame = w / h
|
46 |
+
if aspect_ratio_frame > aspect_ratio_target:
|
47 |
+
new_width = int(h * aspect_ratio_target)
|
48 |
+
x_start = (w - new_width) // 2
|
49 |
+
frame_cropped = frame[:, x_start:x_start + new_width]
|
50 |
+
else:
|
51 |
+
new_height = int(w / aspect_ratio_target)
|
52 |
+
y_start = (h - new_height) // 2
|
53 |
+
frame_cropped = frame[y_start:y_start + new_height, :]
|
54 |
+
frame_resized = cv2.resize(frame_cropped, (target_width, target_height))
|
55 |
+
return frame_resized
|
56 |
+
|
57 |
+
def load_video_to_tensor_with_resize(video_path, target_height=512, target_width=768):
|
58 |
+
cap = cv2.VideoCapture(video_path)
|
59 |
+
frames = []
|
60 |
+
while True:
|
61 |
+
ret, frame = cap.read()
|
62 |
+
if not ret:
|
63 |
+
break
|
64 |
+
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
65 |
+
frame_resized = center_crop_and_resize(frame_rgb, target_height, target_width)
|
66 |
+
frames.append(frame_resized)
|
67 |
+
cap.release()
|
68 |
+
video_np = np.array(frames)
|
69 |
+
video_tensor = torch.tensor(video_np).permute(3, 0, 1, 2).float()
|
70 |
+
video_tensor = (video_tensor / 127.5) - 1.0
|
71 |
+
return video_tensor
|
72 |
+
|
73 |
+
def load_image_to_tensor_with_resize(image_path, target_height=512, target_width=768):
|
74 |
+
image = Image.open(image_path).convert("RGB")
|
75 |
+
image_np = np.array(image)
|
76 |
+
frame_resized = center_crop_and_resize(image_np, target_height, target_width)
|
77 |
+
frame_tensor = torch.tensor(frame_resized).permute(2, 0, 1).float()
|
78 |
+
frame_tensor = (frame_tensor / 127.5) - 1.0
|
79 |
+
# Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
|
80 |
+
return frame_tensor.unsqueeze(0).unsqueeze(2)
|
81 |
+
|
82 |
def main():
|
83 |
+
parser = argparse.ArgumentParser(description='Load models from separate directories and run the pipeline.')
|
84 |
+
|
85 |
+
# Directories
|
86 |
+
parser.add_argument('--ckpt_dir', type=str, required=True,
|
87 |
+
help='Path to the directory containing unet, vae, and scheduler subdirectories')
|
88 |
+
parser.add_argument('--video_path', type=str,
|
89 |
+
help='Path to the input video file (first frame used)')
|
90 |
+
parser.add_argument('--image_path', type=str,
|
91 |
+
help='Path to the input image file')
|
92 |
+
parser.add_argument('--seed', type=int, default="171198")
|
93 |
+
|
94 |
+
# Pipeline parameters
|
95 |
+
parser.add_argument('--num_inference_steps', type=int, default=40, help='Number of inference steps')
|
96 |
+
parser.add_argument('--num_images_per_prompt', type=int, default=1, help='Number of images per prompt')
|
97 |
+
parser.add_argument('--guidance_scale', type=float, default=3, help='Guidance scale for the pipeline')
|
98 |
+
parser.add_argument('--height', type=int, default=512, help='Height of the output video frames')
|
99 |
+
parser.add_argument('--width', type=int, default=768, help='Width of the output video frames')
|
100 |
+
parser.add_argument('--num_frames', type=int, default=121, help='Number of frames to generate in the output video')
|
101 |
+
parser.add_argument('--frame_rate', type=int, default=25, help='Frame rate for the output video')
|
102 |
+
|
103 |
+
# Prompts
|
104 |
+
parser.add_argument('--prompt', type=str,
|
105 |
+
default='A man wearing a black leather jacket and blue jeans is riding a Harley Davidson motorcycle down a paved road. The man has short brown hair and is wearing a black helmet. The motorcycle is a dark red color with a large front fairing. The road is surrounded by green grass and trees. There is a gas station on the left side of the road with a red and white sign that says "Oil" and "Diner".',
|
106 |
+
help='Text prompt to guide generation')
|
107 |
+
parser.add_argument('--negative_prompt', type=str,
|
108 |
+
default='worst quality, inconsistent motion, blurry, jittery, distorted',
|
109 |
+
help='Negative prompt for undesired features')
|
110 |
+
|
111 |
args = parser.parse_args()
|
112 |
|
113 |
# Paths for the separate mode directories
|
114 |
+
ckpt_dir = Path(args.ckpt_dir)
|
115 |
+
unet_dir = ckpt_dir / 'unet'
|
116 |
+
vae_dir = ckpt_dir / 'vae'
|
117 |
+
scheduler_dir = ckpt_dir / 'scheduler'
|
118 |
|
119 |
# Load models
|
120 |
vae = load_vae(vae_dir)
|
121 |
unet = load_unet(unet_dir)
|
122 |
scheduler = load_scheduler(scheduler_dir)
|
|
|
|
|
123 |
patchifier = SymmetricPatchifier(patch_size=1)
|
124 |
+
text_encoder = T5EncoderModel.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder").to(
|
125 |
+
"cuda")
|
126 |
+
tokenizer = T5Tokenizer.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer")
|
127 |
|
128 |
# Use submodels for the pipeline
|
129 |
submodel_dict = {
|
130 |
+
"transformer": unet,
|
131 |
"patchifier": patchifier,
|
132 |
+
"text_encoder": text_encoder,
|
133 |
+
"tokenizer": tokenizer,
|
134 |
"scheduler": scheduler,
|
135 |
"vae": vae,
|
136 |
}
|
137 |
|
138 |
+
pipeline = VideoPixArtAlphaPipeline(**submodel_dict).to("cuda")
|
139 |
+
|
140 |
+
# Load media (video or image)
|
141 |
+
if args.video_path:
|
142 |
+
media_items = load_video_to_tensor_with_resize(args.video_path, args.height, args.width).unsqueeze(0)
|
143 |
+
elif args.image_path:
|
144 |
+
media_items = load_image_to_tensor_with_resize(args.image_path, args.height, args.width)
|
145 |
+
else:
|
146 |
+
raise ValueError("Either --video_path or --image_path must be provided.")
|
147 |
+
|
148 |
+
# Prepare input for the pipeline
|
149 |
+
sample = {
|
150 |
+
"prompt": args.prompt,
|
151 |
+
'prompt_attention_mask': None,
|
152 |
+
'negative_prompt': args.negative_prompt,
|
153 |
+
'negative_prompt_attention_mask': None,
|
154 |
+
'media_items': media_items,
|
155 |
+
}
|
156 |
|
157 |
+
generator = torch.Generator(device="cpu").manual_seed(args.seed)
|
158 |
|
159 |
+
# Run the pipeline
|
160 |
images = pipeline(
|
161 |
+
num_inference_steps=args.num_inference_steps,
|
162 |
+
num_images_per_prompt=args.num_images_per_prompt,
|
163 |
+
guidance_scale=args.guidance_scale,
|
164 |
+
generator=generator,
|
165 |
output_type="pt",
|
166 |
callback_on_step_end=None,
|
167 |
+
height=args.height,
|
168 |
+
width=args.width,
|
169 |
+
num_frames=args.num_frames,
|
170 |
+
frame_rate=args.frame_rate,
|
171 |
**sample,
|
172 |
is_video=True,
|
173 |
vae_per_channel_normalize=True,
|
174 |
+
conditioning_method=ConditioningMethod.FIRST_FRAME
|
175 |
).images
|
176 |
|
177 |
+
# Save output video
|
178 |
+
for i in range(images.shape[0]):
|
179 |
+
video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
|
180 |
+
video_np = (video_np * 255).astype(np.uint8)
|
181 |
+
fps = args.frame_rate
|
182 |
+
height, width = video_np.shape[1:3]
|
183 |
+
filename = lambda base, ext, dir='.': next(
|
184 |
+
os.path.join(dir, f"{base}_{i}{ext}") for i in range(1000) if
|
185 |
+
not os.path.exists(os.path.join(dir, f"{base}_{i}{ext}")))
|
186 |
+
out = cv2.VideoWriter(filename(f"video_output_{i}", ".mp4", "."), cv2.VideoWriter_fourcc(*'mp4v'), fps,
|
187 |
+
(width, height))
|
188 |
+
for frame in video_np[..., ::-1]:
|
189 |
+
out.write(frame)
|
190 |
+
out.release()
|
191 |
+
|
192 |
|
193 |
if __name__ == "__main__":
|
194 |
main()
|