| | import torch |
| | from torch.utils.data import Dataset |
| | import torchvision.transforms as T |
| | import os |
| | import random |
| | import numpy as np |
| |
|
| | from PIL import Image, ImageDraw |
| |
|
| | from datasets import load_dataset |
| |
|
| | from .trainer import OminiModel, get_config, train |
| | from ..pipeline.flux_omini import Condition, convert_to_condition, generate |
| |
|
| |
|
| | class ImageConditionDataset(Dataset): |
| | def __init__( |
| | self, |
| | base_dataset, |
| | condition_size=(512, 512), |
| | target_size=(512, 512), |
| | condition_type: str = "canny", |
| | drop_text_prob: float = 0.1, |
| | drop_image_prob: float = 0.1, |
| | return_pil_image: bool = False, |
| | position_scale=1.0, |
| | ): |
| | self.base_dataset = base_dataset |
| | self.condition_size = condition_size |
| | self.target_size = target_size |
| | self.condition_type = condition_type |
| | self.drop_text_prob = drop_text_prob |
| | self.drop_image_prob = drop_image_prob |
| | self.return_pil_image = return_pil_image |
| | self.position_scale = position_scale |
| |
|
| | self.to_tensor = T.ToTensor() |
| |
|
| | def __len__(self): |
| | return len(self.base_dataset) |
| |
|
| | def __get_condition__(self, image, condition_type): |
| | condition_size = self.condition_size |
| | position_delta = np.array([0, 0]) |
| | if condition_type in ["canny", "coloring", "deblurring", "depth"]: |
| | image, kwargs = image.resize(condition_size), {} |
| | if condition_type == "deblurring": |
| | blur_radius = random.randint(1, 10) |
| | kwargs["blur_radius"] = blur_radius |
| | condition_img = convert_to_condition(condition_type, image, **kwargs) |
| | elif condition_type == "depth_pred": |
| | depth_img = convert_to_condition("depth", image) |
| | condition_img = image.resize(condition_size) |
| | image = depth_img.resize(condition_size) |
| | elif condition_type == "fill": |
| | condition_img = image.resize(condition_size).convert("RGB") |
| | w, h = image.size |
| | x1, x2 = sorted([random.randint(0, w), random.randint(0, w)]) |
| | y1, y2 = sorted([random.randint(0, h), random.randint(0, h)]) |
| | mask = Image.new("L", image.size, 0) |
| | draw = ImageDraw.Draw(mask) |
| | draw.rectangle([x1, y1, x2, y2], fill=255) |
| | if random.random() > 0.5: |
| | mask = Image.eval(mask, lambda a: 255 - a) |
| | condition_img = Image.composite( |
| | image, Image.new("RGB", image.size, (0, 0, 0)), mask |
| | ) |
| | elif condition_type == "sr": |
| | condition_img = image.resize(condition_size) |
| | position_delta = np.array([0, -condition_size[0] // 16]) |
| | else: |
| | raise ValueError(f"Condition type {condition_type} is not implemented.") |
| | return condition_img, position_delta |
| |
|
| | def __getitem__(self, idx): |
| | image = self.base_dataset[idx]["jpg"] |
| | image = image.resize(self.target_size).convert("RGB") |
| | description = self.base_dataset[idx]["json"]["prompt"] |
| |
|
| | condition_size = self.condition_size |
| | position_scale = self.position_scale |
| |
|
| | condition_img, position_delta = self.__get_condition__( |
| | image, self.condition_type |
| | ) |
| |
|
| | |
| | drop_text = random.random() < self.drop_text_prob |
| | drop_image = random.random() < self.drop_image_prob |
| |
|
| | if drop_text: |
| | description = "" |
| | if drop_image: |
| | condition_img = Image.new("RGB", condition_size, (0, 0, 0)) |
| |
|
| | return { |
| | "image": self.to_tensor(image), |
| | "condition_0": self.to_tensor(condition_img), |
| | "condition_type_0": self.condition_type, |
| | "position_delta_0": position_delta, |
| | "description": description, |
| | **({"pil_image": [image, condition_img]} if self.return_pil_image else {}), |
| | **({"position_scale_0": position_scale} if position_scale != 1.0 else {}), |
| | } |
| |
|
| |
|
| | @torch.no_grad() |
| | def test_function(model, save_path, file_name): |
| | condition_size = model.training_config["dataset"]["condition_size"] |
| | target_size = model.training_config["dataset"]["target_size"] |
| |
|
| | position_delta = model.training_config["dataset"].get("position_delta", [0, 0]) |
| | position_scale = model.training_config["dataset"].get("position_scale", 1.0) |
| |
|
| | adapter = model.adapter_names[2] |
| | condition_type = model.training_config["condition_type"] |
| | test_list = [] |
| |
|
| | if condition_type in ["canny", "coloring", "deblurring", "depth"]: |
| | image = Image.open("assets/vase_hq.jpg") |
| | image = image.resize(condition_size) |
| | condition_img = convert_to_condition(condition_type, image, 5) |
| | condition = Condition(condition_img, adapter, position_delta, position_scale) |
| | test_list.append((condition, "A beautiful vase on a table.")) |
| | elif condition_type == "depth_pred": |
| | image = Image.open("assets/vase_hq.jpg") |
| | image = image.resize(condition_size) |
| | condition = Condition(image, adapter, position_delta, position_scale) |
| | test_list.append((condition, "A beautiful vase on a table.")) |
| | elif condition_type == "fill": |
| | condition_img = ( |
| | Image.open("./assets/vase_hq.jpg").resize(condition_size).convert("RGB") |
| | ) |
| | mask = Image.new("L", condition_img.size, 0) |
| | draw = ImageDraw.Draw(mask) |
| | a = condition_img.size[0] // 4 |
| | b = a * 3 |
| | draw.rectangle([a, a, b, b], fill=255) |
| | condition_img = Image.composite( |
| | condition_img, Image.new("RGB", condition_img.size, (0, 0, 0)), mask |
| | ) |
| | condition = Condition(condition, adapter, position_delta, position_scale) |
| | test_list.append((condition, "A beautiful vase on a table.")) |
| | elif condition_type == "super_resolution": |
| | image = Image.open("assets/vase_hq.jpg") |
| | image = image.resize(condition_size) |
| | condition = Condition(image, adapter, position_delta, position_scale) |
| | test_list.append((condition, "A beautiful vase on a table.")) |
| | else: |
| | raise NotImplementedError |
| | os.makedirs(save_path, exist_ok=True) |
| | for i, (condition, prompt) in enumerate(test_list): |
| | generator = torch.Generator(device=model.device) |
| | generator.manual_seed(42) |
| |
|
| | res = generate( |
| | model.flux_pipe, |
| | prompt=prompt, |
| | conditions=[condition], |
| | height=target_size[1], |
| | width=target_size[0], |
| | generator=generator, |
| | model_config=model.model_config, |
| | kv_cache=model.model_config.get("independent_condition", False), |
| | ) |
| | file_path = os.path.join(save_path, f"{file_name}_{condition_type}_{i}.jpg") |
| | res.images[0].save(file_path) |
| |
|
| |
|
| | def main(): |
| | |
| | config = get_config() |
| | training_config = config["train"] |
| | torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) |
| |
|
| | |
| | dataset = load_dataset( |
| | "webdataset", |
| | data_files={"train": training_config["dataset"]["urls"]}, |
| | split="train", |
| | cache_dir="cache/t2i2m", |
| | num_proc=32, |
| | ) |
| |
|
| | |
| | dataset = ImageConditionDataset( |
| | dataset, |
| | condition_size=training_config["dataset"]["condition_size"], |
| | target_size=training_config["dataset"]["target_size"], |
| | condition_type=training_config["condition_type"], |
| | drop_text_prob=training_config["dataset"]["drop_text_prob"], |
| | drop_image_prob=training_config["dataset"]["drop_image_prob"], |
| | position_scale=training_config["dataset"].get("position_scale", 1.0), |
| | ) |
| |
|
| | |
| | trainable_model = OminiModel( |
| | flux_pipe_id=config["flux_path"], |
| | lora_config=training_config["lora_config"], |
| | device=f"cuda", |
| | dtype=getattr(torch, config["dtype"]), |
| | optimizer_config=training_config["optimizer"], |
| | model_config=config.get("model", {}), |
| | gradient_checkpointing=training_config.get("gradient_checkpointing", False), |
| | ) |
| |
|
| | train(dataset, trainable_model, config, test_function) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | main() |
| |
|
