Spaces:
Runtime error
Runtime error
File size: 9,366 Bytes
50edba0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 |
import os
import math
import time
import numpy as np
import random
import threading
from PIL import Image, ImageOps
from moviepy.editor import VideoFileClip
from datetime import datetime, timedelta
from huggingface_hub import hf_hub_download, snapshot_download
import insightface
from insightface.app import FaceAnalysis
from facexlib.parsing import init_parsing_model
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
import torch
from diffusers import CogVideoXDPMScheduler
from diffusers.utils import load_image
from diffusers.image_processor import VaeImageProcessor
from diffusers.training_utils import free_memory
from util.utils import *
from util.rife_model import load_rife_model, rife_inference_with_latents
from models.utils import process_face_embeddings
from models.transformer_consisid import ConsisIDTransformer3DModel
from models.pipeline_consisid import ConsisIDPipeline
from models.eva_clip import create_model_and_transforms
from models.eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
from models.eva_clip.utils_qformer import resize_numpy_image_long
import argparse
device = "cuda" if torch.cuda.is_available() else "cpu"
def main():
parser = argparse.ArgumentParser(description="ConsisID Command Line Interface")
parser.add_argument("image_path", type=str, help="Path to the input image")
parser.add_argument("prompt", type=str, help="Prompt text for the generation")
parser.add_argument("--num_inference_steps", type=int, default=50, help="Number of inference steps")
parser.add_argument("--guidance_scale", type=float, default=7.0, help="Guidance scale")
parser.add_argument("--seed", type=int, default=42, help="Random seed for generation")
parser.add_argument("--output_dir", type=str, default="./output", help="Directory to save the output video")
parser.add_argument("--num_videos", type=int, default=1, help="Number of videos to generate")
args = parser.parse_args()
# Download models
hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran")
snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife")
snapshot_download(repo_id="BestWishYsh/ConsisID-preview", local_dir="BestWishYsh/ConsisID-preview")
model_path = "BestWishYsh/ConsisID-preview"
lora_path = None
lora_rank = 128
dtype = torch.bfloat16
if os.path.exists(os.path.join(model_path, "transformer_ema")):
subfolder = "transformer_ema"
else:
subfolder = "transformer"
transformer = ConsisIDTransformer3DModel.from_pretrained_cus(model_path, subfolder=subfolder)
scheduler = CogVideoXDPMScheduler.from_pretrained(model_path, subfolder="scheduler")
try:
is_kps = transformer.config.is_kps
except:
is_kps = False
# 1. load face helper models
face_helper = FaceRestoreHelper(
upscale_factor=1,
face_size=512,
crop_ratio=(1, 1),
det_model='retinaface_resnet50',
save_ext='png',
device=device,
model_rootpath=os.path.join(model_path, "face_encoder")
)
face_helper.face_parse = None
face_helper.face_parse = init_parsing_model(model_name='bisenet', device=device, model_rootpath=os.path.join(model_path, "face_encoder"))
face_helper.face_det.eval()
face_helper.face_parse.eval()
model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', os.path.join(model_path, "face_encoder", "EVA02_CLIP_L_336_psz14_s6B.pt"), force_custom_clip=True)
face_clip_model = model.visual
face_clip_model.eval()
eva_transform_mean = getattr(face_clip_model, 'image_mean', OPENAI_DATASET_MEAN)
eva_transform_std = getattr(face_clip_model, 'image_std', OPENAI_DATASET_STD)
if not isinstance(eva_transform_mean, (list, tuple)):
eva_transform_mean = (eva_transform_mean,) * 3
if not isinstance(eva_transform_std, (list, tuple)):
eva_transform_std = (eva_transform_std,) * 3
eva_transform_mean = eva_transform_mean
eva_transform_std = eva_transform_std
face_main_model = FaceAnalysis(name='antelopev2', root=os.path.join(model_path, "face_encoder"), providers=['CUDAExecutionProvider'])
handler_ante = insightface.model_zoo.get_model(f'{model_path}/face_encoder/models/antelopev2/glintr100.onnx', providers=['CUDAExecutionProvider'])
face_main_model.prepare(ctx_id=0, det_size=(640, 640))
handler_ante.prepare(ctx_id=0)
face_clip_model.to(device, dtype=dtype)
face_helper.face_det.to(device)
face_helper.face_parse.to(device)
transformer.to(device, dtype=dtype)
free_memory()
pipe = ConsisIDPipeline.from_pretrained(model_path, transformer=transformer, scheduler=scheduler, torch_dtype=dtype)
# If you're using with lora, add this code
if lora_path:
pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1")
pipe.fuse_lora(lora_scale=1 / lora_rank)
scheduler_args = {}
if "variance_type" in pipe.scheduler.config:
variance_type = pipe.scheduler.config.variance_type
if variance_type in ["learned", "learned_range"]:
variance_type = "fixed_small"
scheduler_args["variance_type"] = variance_type
pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)
#pipe.to(device)
pipe.enable_model_cpu_offload()
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()
os.makedirs(args.output_dir, exist_ok=True)
upscale_model = load_sd_upscale("model_real_esran/RealESRGAN_x4.pth", device)
frame_interpolation_model = load_rife_model("model_rife")
def infer(
prompt: str,
image_input: str,
num_inference_steps: int,
guidance_scale: float,
seed: int = 42,
):
if seed == -1:
seed = random.randint(0, 2**8 - 1)
id_image = np.array(ImageOps.exif_transpose(Image.open(image_input)).convert("RGB"))
id_image = resize_numpy_image_long(id_image, 1024)
id_cond, id_vit_hidden, align_crop_face_image, face_kps = process_face_embeddings(face_helper, face_clip_model, handler_ante,
eva_transform_mean, eva_transform_std,
face_main_model, device, dtype, id_image,
original_id_image=id_image, is_align_face=True,
cal_uncond=False)
if is_kps:
kps_cond = face_kps
else:
kps_cond = None
tensor = align_crop_face_image.cpu().detach()
tensor = tensor.squeeze()
tensor = tensor.permute(1, 2, 0)
tensor = tensor.numpy() * 255
tensor = tensor.astype(np.uint8)
image = ImageOps.exif_transpose(Image.fromarray(tensor))
prompt = prompt.strip('"')
generator = torch.Generator(device).manual_seed(seed) if seed else None
video_pt = pipe(
prompt=prompt,
image=image,
num_videos_per_prompt=1,
num_inference_steps=num_inference_steps,
num_frames=49,
use_dynamic_cfg=False,
guidance_scale=guidance_scale,
generator=generator,
id_vit_hidden=id_vit_hidden,
id_cond=id_cond,
kps_cond=kps_cond,
output_type="pt",
).frames
free_memory()
return (video_pt, seed)
def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8, output_dir = "output"):
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
video_path = f"./{output_dir}/{timestamp}.mp4"
os.makedirs(os.path.dirname(video_path), exist_ok=True)
export_to_video(tensor, video_path, fps=fps)
return video_path
def convert_to_gif(video_path):
clip = VideoFileClip(video_path)
gif_path = video_path.replace(".mp4", ".gif")
clip.write_gif(gif_path, fps=8)
return gif_path
for i in range(args.num_videos):
seed = random.randint(0, 2**8 - 1) if args.seed == -1 else args.seed + i
latents, seed = infer(
args.prompt,
args.image_path,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=seed,
)
batch_size = latents.shape[0]
batch_video_frames = []
for batch_idx in range(batch_size):
pt_image = latents[batch_idx]
pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])])
image_np = VaeImageProcessor.pt_to_numpy(pt_image)
image_pil = VaeImageProcessor.numpy_to_pil(image_np)
batch_video_frames.append(image_pil)
video_path = save_video(batch_video_frames[0], fps=math.ceil((len(batch_video_frames[0]) - 1) / 6), output_dir=args.output_dir)
gif_path = convert_to_gif(video_path)
print(f"Video {i+1} saved to: {video_path}")
print(f"GIF {i+1} saved to: {gif_path}")
if __name__ == "__main__":
main() |