# Prediction interface for Cog ⚙️
# https://cog.run/python

from cog import BasePredictor, Input, Path
from typing import List
from modelscope.pipelines import pipeline
from util import save_images, read_img_result

img_save_folder = "SaveImages"
params = {
    "show_debug": True,
    "image_count": 2,
    "ddim_steps": 20,
}


class Predictor(BasePredictor):
    def setup(self) -> None:
        self.model = pipeline('my-anytext-task',
                              model='damo/cv_anytext_text_generation_editing', model_revision='v1.1.3')

    def predict(
        self,
        mode: str = Input(description="Select model type", default="text-generation", choices=[
            "text-generation", "text-editing"]),
        prompt: str = Input(description="Input prompt",
                            default='photo of caramel macchiato coffee on the table, top-down perspective, with "Any" "Text" written on it using cream'),
        seed: int = Input(description="", default=66273235,
                          ge=0, le=66273235),
        draw_pos: Path = Input(
            description="", default='https://replicate.delivery/pbxt/WwkQofrF9CQuJ6v3slDfyha3YB1teWJpnWLBsVSdvfeFh2tRC/image_0.png'),
        ori_image: Path = Input(
            description="", default='https://replicate.delivery/pbxt/0eNDe5B73njlkkSJfJeZ4Riiww3OYci679bnbaf0cNzLh2tRC/image_1.png'),
        use_fp32: bool = Input(description="", default=False),
        no_translator: bool = Input(description="", default=False),
    ) -> List[str]:
        input_data = {}
        if mode == "text-generation":
            input_data = {
                "prompt": prompt,
                "seed": seed,
                "draw_pos": str(draw_pos),
                "use_fp32": use_fp32,
                "no_translator": no_translator,
            }
        else:
            input_data = {
                "prompt": prompt,
                "seed": seed,
                "draw_pos": str(draw_pos),
                "ori_image": str(ori_image),
                "use_fp32": use_fp32,
                "no_translator": no_translator,
            }
        print("input-data:", input_data)
        results, rtn_code, rtn_warning, debug_info = self.model(
            input_data, mode=mode, **params)
        if rtn_warning:
            print(rtn_warning)
            return []
        res = []
        if rtn_code >= 0:
            res = read_img_result(results)
            # save_images(results, img_save_folder)
            print(f'Done, result images are saved in: {img_save_folder}')
        return res
