init commit
This commit is contained in:
root
+260
@@ -0,0 +1,260 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from diffusers import (EulerDiscreteScheduler, EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler, PNDMScheduler, DDIMScheduler)
|
||||
from omegaconf import OmegaConf
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
from safetensors.torch import load_file as load_safetensors
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from ruyi.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
|
||||
from ruyi.models.autoencoder_magvit import AutoencoderKLMagvit
|
||||
from ruyi.models.transformer3d import HunyuanTransformer3DModel
|
||||
from ruyi.pipeline.pipeline_ruyi_inpaint import RuyiInpaintPipeline
|
||||
from ruyi.utils.lora_utils import merge_lora, unmerge_lora
|
||||
from ruyi.utils.utils import get_image_to_video_latent, save_videos_grid
|
||||
|
||||
# Input and output
|
||||
start_image_path = "assets/girl_01.jpg"
|
||||
end_image_path = "assets/girl_02.jpg" # Can be None for start-image-to-video
|
||||
output_video_path = "outputs/example_01.mp4"
|
||||
|
||||
# Video settings
|
||||
video_length = 120 # The max video length is 120 frames (24 frames per second)
|
||||
base_resolution = 512 # # The pixels in the generated video are approximately 512 x 512. Values in the range of [384, 896] typically produce good video quality.
|
||||
video_size = None # Override base_resolution. Format: [height, width], e.g., [384, 672]
|
||||
# Control settings
|
||||
aspect_ratio = "16:9" # Choose in ["16:9", "9:16"], note that this is only the hint
|
||||
motion = "auto" # Motion control, choose in ["1", "2", "3", "4", "auto"]
|
||||
camera_direction = "auto" # Camera control, choose in ["static", "left", "right", "up", "down", "auto"]
|
||||
# Sampler settings
|
||||
steps = 25
|
||||
cfg = 7.0
|
||||
scheduler_name = "DDIM" # Choose in ["Euler", "Euler A", "DPM++", "PNDM","DDIM"]
|
||||
|
||||
# GPU memory settings
|
||||
low_gpu_memory_mode = False # Low gpu memory mode
|
||||
gpu_offload_steps = 5 # Choose in [0, 10, 7, 5, 1], the latter number requires less GPU memory but longer time
|
||||
|
||||
# Random seed
|
||||
seed = 42 # The Answer to the Ultimate Question of Life, The Universe, and Everything
|
||||
|
||||
# Model settings
|
||||
config_path = "config/default.yaml"
|
||||
model_name = "Ruyi-Mini-7B"
|
||||
model_type = "Inpaint"
|
||||
model_path = f"models/{model_name}" # (Down)load mode in this path
|
||||
auto_download = True
|
||||
|
||||
# LoRA settings
|
||||
lora_path = None
|
||||
lora_weight = 1.0
|
||||
|
||||
# Other settings
|
||||
weight_dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
|
||||
def get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction):
|
||||
# Default keys
|
||||
p_default_key = "p.default"
|
||||
n_default_key = "n.default"
|
||||
|
||||
# Load embeddings
|
||||
if motion == "auto":
|
||||
motion = "0"
|
||||
p_key = f"p.{aspect_ratio.replace(':', 'x')}movie{motion}{camera_direction}"
|
||||
embeddings = pipeline.embeddings
|
||||
|
||||
# Get embeddings
|
||||
positive_embeds = embeddings.get(f"{p_key}.emb1", embeddings[f"{p_default_key}.emb1"])
|
||||
positive_attention_mask = embeddings.get(f"{p_key}.mask1", embeddings[f"{p_default_key}.mask1"])
|
||||
positive_embeds_2 = embeddings.get(f"{p_key}.emb2", embeddings[f"{p_default_key}.emb2"])
|
||||
positive_attention_mask_2 = embeddings.get(f"{p_key}.mask2", embeddings[f"{p_default_key}.mask2"])
|
||||
|
||||
negative_embeds = embeddings[f"{n_default_key}.emb1"]
|
||||
negative_attention_mask = embeddings[f"{n_default_key}.mask1"]
|
||||
negative_embeds_2 = embeddings[f"{n_default_key}.emb2"]
|
||||
negative_attention_mask_2 = embeddings[f"{n_default_key}.mask2"]
|
||||
|
||||
return {
|
||||
"positive_embeds": positive_embeds,
|
||||
"positive_attention_mask": positive_attention_mask,
|
||||
"positive_embeds_2": positive_embeds_2,
|
||||
"positive_attention_mask_2": positive_attention_mask_2,
|
||||
|
||||
"negative_embeds": negative_embeds,
|
||||
"negative_attention_mask": negative_attention_mask,
|
||||
"negative_embeds_2": negative_embeds_2,
|
||||
"negative_attention_mask_2": negative_attention_mask_2,
|
||||
}
|
||||
|
||||
|
||||
def try_setup_pipeline(model_path, weight_dtype, config):
|
||||
try:
|
||||
# Get Vae
|
||||
vae = AutoencoderKLMagvit.from_pretrained(
|
||||
model_path,
|
||||
subfolder="vae"
|
||||
).to(weight_dtype)
|
||||
print("Vae loaded ...")
|
||||
|
||||
# Get Transformer
|
||||
transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
|
||||
transformer = HunyuanTransformer3DModel.from_pretrained_2d(
|
||||
model_path,
|
||||
subfolder="transformer",
|
||||
transformer_additional_kwargs=transformer_additional_kwargs
|
||||
).to(weight_dtype)
|
||||
print("Transformer loaded ...")
|
||||
|
||||
# Load Clip
|
||||
clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
model_path, subfolder="image_encoder"
|
||||
).to(weight_dtype)
|
||||
clip_image_processor = CLIPImageProcessor.from_pretrained(
|
||||
model_path, subfolder="image_encoder"
|
||||
)
|
||||
|
||||
# Load sampler and create pipeline
|
||||
Choosen_Scheduler = DDIMScheduler
|
||||
scheduler = Choosen_Scheduler.from_pretrained(
|
||||
model_path,
|
||||
subfolder="scheduler"
|
||||
)
|
||||
pipeline = RuyiInpaintPipeline.from_pretrained(
|
||||
model_path,
|
||||
vae=vae,
|
||||
transformer=transformer,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=weight_dtype,
|
||||
clip_image_encoder=clip_image_encoder,
|
||||
clip_image_processor=clip_image_processor,
|
||||
)
|
||||
|
||||
# Load embeddings
|
||||
embeddings = load_safetensors(os.path.join(model_path, "embeddings.safetensors"))
|
||||
pipeline.embeddings = embeddings
|
||||
print("Pipeline loaded ...")
|
||||
|
||||
return pipeline
|
||||
except Exception as e:
|
||||
print("[Ruyi] Setup pipeline failed:", e)
|
||||
return None
|
||||
|
||||
|
||||
# Init weight_dtype and device
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
# Load config
|
||||
config = OmegaConf.load(config_path)
|
||||
|
||||
# Load images
|
||||
start_img = [Image.open(start_image_path).convert("RGB")]
|
||||
end_img = [Image.open(end_image_path).convert("RGB")] if end_image_path is not None else None
|
||||
|
||||
# Init model
|
||||
repo_id = f"IamCreateAI/{model_name}"
|
||||
pipeline = try_setup_pipeline(model_path, weight_dtype, config)
|
||||
if pipeline is None and auto_download:
|
||||
print(f"Downloading {model_name} ...")
|
||||
|
||||
# Download the model
|
||||
snapshot_download(repo_id=repo_id, local_dir=model_path)
|
||||
|
||||
pipeline = try_setup_pipeline(model_path, weight_dtype, config)
|
||||
|
||||
if pipeline is None:
|
||||
message = (f"[Load Model Failed] "
|
||||
f"Please download Ruyi model from huggingface repo '{repo_id}', "
|
||||
f"And put it into '{model_path}'.")
|
||||
if not auto_download:
|
||||
message += "\nOr just set auto_download to 'True'."
|
||||
raise FileNotFoundError(message)
|
||||
|
||||
# Setup GPU memory mode
|
||||
if low_gpu_memory_mode:
|
||||
pipeline.enable_sequential_cpu_offload()
|
||||
else:
|
||||
pipeline.enable_model_cpu_offload()
|
||||
|
||||
# Prepare LoRA config
|
||||
loras = {
|
||||
'models': [lora_path] if lora_path is not None else [],
|
||||
'weights': [lora_weight] if lora_path is not None else [],
|
||||
}
|
||||
|
||||
# Count most suitable height and width
|
||||
if video_size is None:
|
||||
aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
|
||||
original_width, original_height = start_img[0].size if type(start_img) is list else Image.open(start_img).size
|
||||
closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size)
|
||||
height, width = [int(x / 16) * 16 for x in closest_size]
|
||||
else:
|
||||
height, width = video_size
|
||||
|
||||
# Set hidden states offload steps
|
||||
if pipeline.model_cpu_offload_flag:
|
||||
pipeline.transformer.hidden_cache_size = gpu_offload_steps
|
||||
else:
|
||||
pipeline.transformer.hidden_cache_size = 0
|
||||
|
||||
# Load Sampler
|
||||
if scheduler_name == "DPM++":
|
||||
noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder='scheduler')
|
||||
elif scheduler_name == "Euler":
|
||||
noise_scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler')
|
||||
elif scheduler_name == "Euler A":
|
||||
noise_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler')
|
||||
elif scheduler_name == "PNDM":
|
||||
noise_scheduler = PNDMScheduler.from_pretrained(model_path, subfolder='scheduler')
|
||||
elif scheduler_name == "DDIM":
|
||||
noise_scheduler = DDIMScheduler.from_pretrained(model_path, subfolder='scheduler')
|
||||
pipeline.scheduler = noise_scheduler
|
||||
|
||||
# Set random seed
|
||||
generator= torch.Generator(device).manual_seed(seed)
|
||||
|
||||
# Load control embeddings
|
||||
embeddings = get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction)
|
||||
|
||||
with torch.no_grad():
|
||||
video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1
|
||||
input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width))
|
||||
|
||||
for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])):
|
||||
pipeline = merge_lora(pipeline, _lora_path, _lora_weight)
|
||||
|
||||
sample = pipeline(
|
||||
prompt_embeds = embeddings["positive_embeds"],
|
||||
prompt_attention_mask = embeddings["positive_attention_mask"],
|
||||
prompt_embeds_2 = embeddings["positive_embeds_2"],
|
||||
prompt_attention_mask_2 = embeddings["positive_attention_mask_2"],
|
||||
|
||||
negative_prompt_embeds = embeddings["negative_embeds"],
|
||||
negative_prompt_attention_mask = embeddings["negative_attention_mask"],
|
||||
negative_prompt_embeds_2 = embeddings["negative_embeds_2"],
|
||||
negative_prompt_attention_mask_2 = embeddings["negative_attention_mask_2"],
|
||||
|
||||
video_length = video_length,
|
||||
height = height,
|
||||
width = width,
|
||||
generator = generator,
|
||||
guidance_scale = cfg,
|
||||
num_inference_steps = steps,
|
||||
|
||||
video = input_video,
|
||||
mask_video = input_video_mask,
|
||||
clip_image = clip_image,
|
||||
).videos
|
||||
|
||||
for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])):
|
||||
pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight)
|
||||
|
||||
# Save the video
|
||||
output_folder = os.path.dirname(output_video_path)
|
||||
if output_folder != '':
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
save_videos_grid(sample, output_video_path, fps=24)
|
||||
Reference in New Issue
Block a user