Spaces:
Runtime error
Runtime error
from __future__ import annotations | |
import os | |
import pathlib | |
import shlex | |
import shutil | |
import subprocess | |
import gradio as gr | |
import PIL.Image | |
import torch | |
def pad_image(image: PIL.Image.Image) -> PIL.Image.Image: | |
w, h = image.size | |
if w == h: | |
return image | |
elif w > h: | |
new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0)) | |
new_image.paste(image, (0, (w - h) // 2)) | |
return new_image | |
else: | |
new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0)) | |
new_image.paste(image, ((h - w) // 2, 0)) | |
return new_image | |
class Trainer: | |
def __init__(self): | |
self.is_running = False | |
self.is_running_message = "Another training is in progress." | |
self.output_dir = pathlib.Path("results") | |
self.data_dir = pathlib.Path("data") | |
self.ref_data_dir = self.data_dir / "ref" | |
self.target_data_dir = self.data_dir / "target" | |
def check_if_running(self) -> dict: | |
if self.is_running: | |
return gr.update(value=self.is_running_message) | |
else: | |
return gr.update(value="No training is running.") | |
def cleanup_dirs(self) -> None: | |
shutil.rmtree(self.output_dir, ignore_errors=True) | |
def prepare_dataset( | |
self, | |
ref_images: list, | |
target_image: PIL.Image, | |
target_mask: PIL.Image, | |
resolution: int | |
) -> None: | |
self.ref_data_dir.mkdir(parents=True) | |
self.target_data_dir.mkdir(parents=True) | |
for i, temp_path in enumerate(ref_images): | |
image = PIL.Image.open(temp_path.name) | |
image = pad_image(image) | |
image = image.resize((resolution, resolution)) | |
image = image.convert("RGB") | |
out_path = self.ref_data_dir / f"{i:03d}.jpg" | |
image.save(out_path, format="JPEG", quality=100) | |
target_image.save(self.target_data_dir / "target.jpg", format="JPEG", quality=100) | |
target_mask.save(self.target_data_dir / "mask.jpg", format="JPEG", quality=100) | |
def run( | |
self, | |
base_model: str, | |
resolution_s: str, | |
n_steps: int, | |
ref_images: list | None, | |
target_image: PIL.Image, | |
target_mask: PIL.Image, | |
unet_learning_rate: float, | |
text_encoder_learning_rate: float, | |
gradient_accumulation: int, | |
fp16: bool, | |
use_8bit_adam: bool, | |
gradient_checkpointing: bool, | |
lora_rank: int, | |
lora_alpha: int, | |
lora_bias: str, | |
lora_dropout: float, | |
) -> tuple[dict, list[pathlib.Path]]: | |
if not torch.cuda.is_available(): | |
raise gr.Error("CUDA is not available.") | |
if self.is_running: | |
return gr.update(value=self.is_running_message), [] | |
if ref_images is None: | |
raise gr.Error("You need to upload reference images.") | |
if target_image is None: | |
raise gr.Error("You need to upload target image.") | |
if target_mask is None: | |
raise gr.Error("You need to upload target mask.") | |
resolution = int(resolution_s) | |
self.cleanup_dirs() | |
self.prepare_dataset(ref_images, target_image, target_mask, resolution) | |
command = f""" | |
accelerate launch train_dreambooth.py \ | |
--pretrained_model_name_or_path={base_model} \ | |
--train_data_dir={self.data_dir} \ | |
--output_dir={self.output_dir} \ | |
--resolution={resolution} \ | |
--gradient_accumulation_steps={gradient_accumulation} \ | |
--unet_learning_rate={unet_learning_rate} \ | |
--text_encoder_learning_rate={text_encoder_learning_rate} \ | |
--max_train_steps={n_steps} \ | |
--train_batch_size=16 \ | |
--lr_scheduler=constant \ | |
--lr_warmup_steps=100 \ | |
--lora_r={lora_rank} \ | |
--lora_alpha={lora_alpha} \ | |
--lora_bias={lora_bias} \ | |
--lora_dropout={lora_dropout} \ | |
""" | |
if fp16: | |
command += " --mixed_precision fp16" | |
if use_8bit_adam: | |
command += " --use_8bit_adam" | |
if gradient_checkpointing: | |
command += " --gradient_checkpointing" | |
with open(self.output_dir / "train.sh", "w") as f: | |
command_s = " ".join(command.split()) | |
f.write(command_s) | |
self.is_running = True | |
res = subprocess.run(shlex.split(command)) | |
self.is_running = False | |
if res.returncode == 0: | |
result_message = "Training Completed!" | |
else: | |
result_message = "Training Failed!" | |
model_paths = sorted(self.output_dir.glob("*")) | |
return gr.update(value=result_message), model_paths | |