agent/sub_tasks.py

import logging

from .data_struct import AnalysisInput, AnalysisOutput
from .data_struct import TransferInput, TransferOutput
from .data_struct import CriteriaInput, CriteriaOutput

from model import QwenUMM
from utils import COLOR_GREEN, COLOR_RESET, extract_json_result, load_image


class AnalysisModule:

    def __init__(self, max_new_tokens: int = 1024):
        self.max_new_tokens = max_new_tokens

    def run(self, model: QwenUMM, task_input: AnalysisInput, logger: logging.Logger) -> AnalysisOutput:
        title = "# ---- Analysis Task ---- #"
        logger.info(title)

        # -- Analysis the style features -- #
        image = {"Picture 1": load_image(task_input.ref_image_or_path)}
        output = model(
            task="txt-gen",
            image=image,
            prompt=f"Analysis Suggestions:\n{task_input.suggestion}",
            sys_prompt=task_input.sys_prompt,
            max_new_tokens=self.max_new_tokens,
        )
        logger.debug(f"Model raw output:\n{output}\n")
        output = extract_json_result(output, logger)
        logger.info(f"{COLOR_GREEN}Extract output: \n{output}{COLOR_RESET}")
        logger.info("# " + "-" * (len(title) - 4) + " #")

        analysis_output = AnalysisOutput()
        setattr(analysis_output, task_input.style_type, output)
        return analysis_output


class TransferModule:

    def __init__(
        self,
        num_inference_steps: int = 25,
        height: int = 1024,
        width: int = 1024,
        seed: int = 42,
    ):
        self.num_inference_steps = num_inference_steps
        self.height = height
        self.width = width
        self.seed = seed

    def run(self, model: QwenUMM, task_input: TransferInput, logger: logging.Logger) -> TransferOutput:
        title = "# ---- Transfer Task ---- #"
        logger.info(title)

        image = {"Picture 1": load_image(task_input.cnt_image_or_path)}
        if task_input.ref_image_or_path is not None:
            image["Picture 2"] = load_image(task_input.ref_image_or_path)
        output = model(
            task="img-gen",
            image=image,
            prompt=task_input.prompt,
            sys_prompt=task_input.sys_prompt,
            output_img_height=self.height,
            output_img_width=self.width,
            num_inference_steps=self.num_inference_steps,
            seed=self.seed,
        )
        logger.debug(f"Model raw output:\n{output}")
        logger.info(f"{COLOR_GREEN} Instruct:\n{task_input.prompt}{COLOR_RESET}")
        logger.info(f"{COLOR_GREEN}Content image:\n{task_input.cnt_image_or_path}{COLOR_RESET}")
        logger.info(f"{COLOR_GREEN}Style image:\n{task_input.ref_image_or_path}{COLOR_RESET}")
        logger.info(f"{COLOR_GREEN}Generate image:\n{output}{COLOR_RESET}")
        logger.info("# " + "-" * (len(title) - 4) + " #")
        transfer_output = TransferOutput(
            instruct=task_input.prompt,
            cnt_image_or_path=task_input.cnt_image_or_path,
            ref_image_or_path=task_input.ref_image_or_path,
            sty_image=output,
        )
        return transfer_output


class CriteriaModule:

    def __init__(self, max_new_tokens: int = 1024):
        self.max_new_tokens = max_new_tokens

    def run(self, model: QwenUMM, task_input: CriteriaInput, logger: logging.Logger) -> CriteriaOutput:
        title = "# ---- Criteria Task ---- #"
        logger.info(title)
        # -- Content preservation score -- #
        subtitle = "| ---- content preservation ---- |"
        logger.info(subtitle)
        cnt_sty_score = model(
            task="txt-gen",
            image={
                "Picture 1": load_image(task_input.cnt_image_or_path),
                "Picture 2": load_image(task_input.sty_image_or_path),
            },
            prompt="",
            sys_prompt=task_input.sys_prompts["cs"],
            max_new_tokens=self.max_new_tokens,
        )
        logger.debug(f"Modal raw output:\n{cnt_sty_score}")
        cnt_sty_score = extract_json_result(cnt_sty_score, logger)
        logger.info(f"{COLOR_GREEN}Extract output:\n{cnt_sty_score}{COLOR_RESET}")
        logger.info("| " + "-" * (len(subtitle) - 4) + " |")

        # -- Style alignment score -- #
        subtitle = "| ---- style alignment ---- |"
        logger.info(subtitle)
        ref_sty_score = model(
            task="txt-gen",
            image={
                "Picture 1": load_image(task_input.ref_image_or_path),
                "Picture 2": load_image(task_input.sty_image_or_path),
            },
            prompt="",
            sys_prompt=task_input.sys_prompts["rs"],
            max_new_tokens=self.max_new_tokens,
        )
        logger.debug(f"Modal raw output:\n{ref_sty_score}")
        ref_sty_score = extract_json_result(ref_sty_score, logger)
        logger.info(f"{COLOR_GREEN}Extract output:\n{ref_sty_score}{COLOR_RESET}")
        logger.info("| " + "-" * (len(subtitle) - 4) + " |")

        # -- Instruct following score
        subtitle = "| ---- instruct following ---- |"
        logger.info(subtitle)
        des_sty_score = model(
            task="txt-gen",
            image={"Picture 1": load_image(task_input.sty_image_or_path)},
            prompt=f"Instruction: {task_input.instruction}",
            sys_prompt=task_input.sys_prompts["ds"],
            max_new_tokens=self.max_new_tokens,
        )
        logger.debug(f"Modal raw output:\n{des_sty_score}")
        des_sty_score = extract_json_result(des_sty_score, logger)
        logger.info(f"{COLOR_GREEN}Extract output:\n{des_sty_score}{COLOR_RESET}")
        logger.info("| " + "-" * (len(subtitle) - 4) + " |")

        logger.info("# " + "-" * (len(title) - 4) + " #")
        criteria_output = CriteriaOutput(
            cs_score=cnt_sty_score,
            rs_score=ref_sty_score,
            ds_score=des_sty_score,
        )
        return criteria_output