# =====================================================
# Apckeyl_RealESRGAN
# Version 13.0
# official_tile_processor.py
# =====================================================

import torch


class OfficialTileProcessor:

    """
    CPU Optimized Version.

    Batch отключён.

    Причина:
    на CPU HuggingFace Batch оказался
    медленнее последовательной обработки.

    Архитектура сохранена.
    При переходе на GPU Batch можно будет
    включить одной функцией.
    """

    def __init__(self, engine):

        self.engine = engine

    # -------------------------------------------------

    def inference(self, tensor):

        with torch.inference_mode():

            return self.engine.infer(
                tensor
            )

    # -------------------------------------------------

    def process(self, tensors):

        outputs = []

        total = len(tensors)

        print()

        print("===== CPU FAST MODE =====")

        print(f"Tiles : {total}")

        print("=========================")

        print()

        for index, tensor in enumerate(
            tensors,
            start=1
        ):

            print(
                f"Tile {index}/{total}"
            )

            outputs.append(

                self.inference(
                    tensor
                )

            )

        return outputs

    # -------------------------------------------------

    def process_tiles(
        self,
        tensors
    ):

        return self.process(
            tensors
        )

    # -------------------------------------------------

    def process_batch(
        self,
        tensors
    ):

        return self.process(
            tensors
        )