Function test_function

omini/train_flux/train_multi_condition.py:57–114

(model, save_path, file_name)

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

    condition_type = model.training_config["condition_type"]
    test_list = []

    condition_list = []
    for i, c_type in enumerate(condition_type):
        if c_type in ["canny", "coloring", "deblurring", "depth"]:
            image = Image.open("assets/vase_hq.jpg")
            image = image.resize(condition_size)
            condition_img = convert_to_condition(c_type, image, 5)
        elif c_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
            )
        else:
            raise NotImplementedError
        condition = Condition(
            condition_img,
            model.adapter_names[i + 2],
            position_delta,
            position_scale,
        )
        condition_list.append(condition)
    test_list.append((condition_list, "A beautiful vase on a table."))
    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_list,
            height=target_size[0],
            width=target_size[1],
            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}_{'|'.join(condition_type)}_{i}.jpg"
        )
        res.images[0].save(file_path)