| 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() |
|
|
