Andres99's picture
Duplicate from Tune-A-Video-library/Tune-A-Video-Training-UI
674b3fa
from __future__ import annotations
import datetime
import os
import pathlib
import shlex
import shutil
import subprocess
import sys
import gradio as gr
import slugify
import torch
from huggingface_hub import HfApi
from omegaconf import OmegaConf
from app_upload import ModelUploader
from utils import save_model_card
sys.path.append('Tune-A-Video')
URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'
ORIGINAL_SPACE_ID = 'Tune-A-Video-library/Tune-A-Video-Training-UI'
SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
class Trainer:
def __init__(self, hf_token: str | None = None):
self.hf_token = hf_token
self.model_uploader = ModelUploader(hf_token)
self.checkpoint_dir = pathlib.Path('checkpoints')
self.checkpoint_dir.mkdir(exist_ok=True)
def download_base_model(self, base_model_id: str) -> str:
model_dir = self.checkpoint_dir / base_model_id
if not model_dir.exists():
org_name = base_model_id.split('/')[0]
org_dir = self.checkpoint_dir / org_name
org_dir.mkdir(exist_ok=True)
subprocess.run(shlex.split(
f'git clone https://huggingface.co/{base_model_id}'),
cwd=org_dir)
return model_dir.as_posix()
def join_model_library_org(self, token: str) -> None:
subprocess.run(
shlex.split(
f'curl -X POST -H "Authorization: Bearer {token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
))
def run(
self,
training_video: str,
training_prompt: str,
output_model_name: str,
overwrite_existing_model: bool,
validation_prompt: str,
base_model: str,
resolution_s: str,
n_steps: int,
learning_rate: float,
gradient_accumulation: int,
seed: int,
fp16: bool,
use_8bit_adam: bool,
checkpointing_steps: int,
validation_epochs: int,
upload_to_hub: bool,
use_private_repo: bool,
delete_existing_repo: bool,
upload_to: str,
remove_gpu_after_training: bool,
input_token: str,
) -> str:
if SPACE_ID == ORIGINAL_SPACE_ID:
raise gr.Error(
'This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU'
)
if not torch.cuda.is_available():
raise gr.Error('CUDA is not available.')
if training_video is None:
raise gr.Error('You need to upload a video.')
if not training_prompt:
raise gr.Error('The training prompt is missing.')
if not validation_prompt:
raise gr.Error('The validation prompt is missing.')
resolution = int(resolution_s)
if not output_model_name:
timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
output_model_name = f'tune-a-video-{timestamp}'
output_model_name = slugify.slugify(output_model_name)
repo_dir = pathlib.Path(__file__).parent
output_dir = repo_dir / 'experiments' / output_model_name
if overwrite_existing_model or upload_to_hub:
shutil.rmtree(output_dir, ignore_errors=True)
output_dir.mkdir(parents=True)
if upload_to_hub:
self.join_model_library_org(
self.hf_token if self.hf_token else input_token)
config = OmegaConf.load('Tune-A-Video/configs/man-surfing.yaml')
config.pretrained_model_path = self.download_base_model(base_model)
config.output_dir = output_dir.as_posix()
config.train_data.video_path = training_video.name # type: ignore
config.train_data.prompt = training_prompt
config.train_data.n_sample_frames = 8
config.train_data.width = resolution
config.train_data.height = resolution
config.train_data.sample_start_idx = 0
config.train_data.sample_frame_rate = 1
config.validation_data.prompts = [validation_prompt]
config.validation_data.video_length = 8
config.validation_data.width = resolution
config.validation_data.height = resolution
config.validation_data.num_inference_steps = 50
config.validation_data.guidance_scale = 7.5
config.learning_rate = learning_rate
config.gradient_accumulation_steps = gradient_accumulation
config.train_batch_size = 1
config.max_train_steps = n_steps
config.checkpointing_steps = checkpointing_steps
config.validation_steps = validation_epochs
config.seed = seed
config.mixed_precision = 'fp16' if fp16 else ''
config.use_8bit_adam = use_8bit_adam
config_path = output_dir / 'config.yaml'
with open(config_path, 'w') as f:
OmegaConf.save(config, f)
command = f'accelerate launch Tune-A-Video/train_tuneavideo.py --config {config_path}'
subprocess.run(shlex.split(command))
save_model_card(save_dir=output_dir,
base_model=base_model,
training_prompt=training_prompt,
test_prompt=validation_prompt,
test_image_dir='samples')
message = 'Training completed!'
print(message)
if upload_to_hub:
upload_message = self.model_uploader.upload_model(
folder_path=output_dir.as_posix(),
repo_name=output_model_name,
upload_to=upload_to,
private=use_private_repo,
delete_existing_repo=delete_existing_repo,
input_token=input_token)
print(upload_message)
message = message + '\n' + upload_message
if remove_gpu_after_training:
space_id = os.getenv('SPACE_ID')
if space_id:
api = HfApi(
token=self.hf_token if self.hf_token else input_token)
api.request_space_hardware(repo_id=space_id,
hardware='cpu-basic')
return message