File size: 5,892 Bytes
59aed9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | 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
|