predict.py

"""
Prediction interface for Cog ⚙️
https://github.com/replicate/cog/blob/main/docs/python.md
"""

import copy
import os
from typing import List
import numpy as np

from PIL import Image
from cog import BasePredictor, BaseModel, Input, Path
from fooocusapi.utils.lora_manager import LoraManager
from fooocusapi.utils.file_utils import output_dir
from fooocusapi.models.common.task import GenerationFinishReason
from fooocusapi.configs.default import (
    available_aspect_ratios,
    uov_methods,
    outpaint_expansions,
    default_styles,
    default_base_model_name,
    default_refiner_model_name,
    default_loras,
    default_refiner_switch,
    default_cfg_scale,
    default_prompt_negative
)

from fooocusapi.parameters import ImageGenerationParams
from fooocusapi.task_queue import TaskType


class Output(BaseModel):
    """
    Output model
    """
    seeds: List[str]
    paths: List[Path]


class Predictor(BasePredictor):
    """Predictor"""
    def setup(self) -> None:
        """
        Load the model into memory to make running multiple predictions efficient
        """
        from main import pre_setup
        pre_setup()

    def predict(
        self,
        prompt: str = Input(
            default='',
            description="Prompt for image generation"),
        negative_prompt: str = Input(
            default=default_prompt_negative,
            description="Negative prompt for image generation"),
        style_selections: str = Input(
            default=','.join(default_styles),
            description="Fooocus styles applied for image generation, separated by comma"),
        performance_selection: str = Input(
            default='Speed', 
            choices=['Speed', 'Quality', 'Extreme Speed', 'Lightning'],
            description="Performance selection"),
        aspect_ratios_selection: str = Input(
            default='1152*896',
            choices=available_aspect_ratios,
            description="The generated image's size"),
        image_number: int = Input(
            default=1,
            ge=1, le=8,
            description="How many image to generate"),
        image_seed: int = Input(
            default=-1,
            description="Seed to generate image, -1 for random"),
        use_default_loras: bool = Input(
            default=True,
            description="Use default LoRAs"),
        loras_custom_urls: str = Input(
            default="",
            description="Custom LoRAs URLs in the format 'url,weight' provide multiple separated by ; (example 'url1,0.3;url2,0.1')"),
        sharpness: float = Input(
            default=2.0,
            ge=0.0, le=30.0),
        guidance_scale: float = Input(
            default=default_cfg_scale,
            ge=1.0, le=30.0),
        refiner_switch: float = Input(
            default=default_refiner_switch,
            ge=0.1, le=1.0),
        uov_input_image: Path = Input(
            default=None,
            description="Input image for upscale or variation, keep None for not upscale or variation"),
        uov_method: str = Input(
            default='Disabled',
            choices=uov_methods),
        uov_upscale_value: float = Input(
            default=0,
            description="Only when Upscale (Custom)"),
        inpaint_additional_prompt: str = Input(
            default='',
            description="Prompt for image generation"),
        inpaint_input_image: Path = Input(
            default=None,
            description="Input image for inpaint or outpaint, keep None for not inpaint or outpaint. Please noticed, `uov_input_image` has bigger priority is not None."),
        inpaint_input_mask: Path = Input(
            default=None,
            description="Input mask for inpaint"),
        outpaint_selections: str = Input(
            default='',
            description="Outpaint expansion selections, literal 'Left', 'Right', 'Top', 'Bottom' separated by comma"),
        outpaint_distance_left: int = Input(
            default=0,
            description="Outpaint expansion distance from Left of the image"),
        outpaint_distance_top: int = Input(
            default=0,
            description="Outpaint expansion distance from Top of the image"),
        outpaint_distance_right: int = Input(
            default=0,
            description="Outpaint expansion distance from Right of the image"),
        outpaint_distance_bottom: int = Input(
            default=0,
            description="Outpaint expansion distance from Bottom of the image"),
        cn_img1: Path = Input(
            default=None,
            description="Input image for image prompt. If all cn_img[n] are None, image prompt will not applied."),
        cn_stop1: float = Input(
            default=None,
            ge=0, le=1,
            description="Stop at for image prompt, None for default value"),
        cn_weight1: float = Input(
            default=None,
            ge=0, le=2,
            description="Weight for image prompt, None for default value"),
        cn_type1: str = Input(
            default='ImagePrompt',
            choices=['ImagePrompt', 'FaceSwap', 'PyraCanny', 'CPDS'],
            description="ControlNet type for image prompt"),
        cn_img2: Path = Input(
            default=None,
            description="Input image for image prompt. If all cn_img[n] are None, image prompt will not applied."),
        cn_stop2: float = Input(
            default=None,
            ge=0, le=1,
            description="Stop at for image prompt, None for default value"),
        cn_weight2: float = Input(
            default=None,
            ge=0, le=2,
            description="Weight for image prompt, None for default value"),
        cn_type2: str = Input(
            default='ImagePrompt',
            choices=['ImagePrompt', 'FaceSwap', 'PyraCanny', 'CPDS'],
            description="ControlNet type for image prompt"),
        cn_img3: Path = Input(
            default=None,
            description="Input image for image prompt. If all cn_img[n] are None, image prompt will not applied."),
        cn_stop3: float = Input(
            default=None,
            ge=0, le=1,
            description="Stop at for image prompt, None for default value"),
        cn_weight3: float = Input(
            default=None,
            ge=0, le=2,
            description="Weight for image prompt, None for default value"),
        cn_type3: str = Input(
            default='ImagePrompt',
            choices=['ImagePrompt', 'FaceSwap', 'PyraCanny', 'CPDS'],
            description="ControlNet type for image prompt"),
        cn_img4: Path = Input(
            default=None,
            description="Input image for image prompt. If all cn_img[n] are None, image prompt will not applied."),
        cn_stop4: float = Input(
            default=None,
            ge=0, le=1,
            description="Stop at for image prompt, None for default value"),
        cn_weight4: float = Input(
            default=None,
            ge=0, le=2,
            description="Weight for image prompt, None for default value"),
        cn_type4: str = Input(
            default='ImagePrompt',
            choices=['ImagePrompt', 'FaceSwap', 'PyraCanny', 'CPDS'],
            description="ControlNet type for image prompt")
    ) -> Output:
        """Run a single prediction on the model"""
        from modules import flags
        from modules.sdxl_styles import legal_style_names
        from fooocusapi.worker import blocking_get_task_result, worker_queue

        base_model_name = default_base_model_name
        refiner_model_name = default_refiner_model_name

        lora_manager = LoraManager()
        
        # Use default loras if selected
        loras = copy.copy(default_loras) if use_default_loras else []

        # add custom user loras if provided
        if loras_custom_urls:
            urls = [url.strip() for url in loras_custom_urls.split(';')]

            loras_with_weights = [url.split(',') for url in urls]

            custom_lora_paths = lora_manager.check([lw[0] for lw in loras_with_weights])
            custom_loras = [[path, float(lw[1]) if len(lw) > 1 else 1.0] for path, lw in
                            zip(custom_lora_paths, loras_with_weights)]

            loras.extend(custom_loras)

        style_selections_arr = []
        for s in style_selections.strip().split(','):
            style = s.strip()
            if style in legal_style_names:
                style_selections_arr.append(style)

        if uov_input_image is not None:
            im = Image.open(str(uov_input_image))
            uov_input_image = np.array(im)

        inpaint_input_image_dict = None
        if inpaint_input_image is not None:
            im = Image.open(str(inpaint_input_image))
            inpaint_input_image = np.array(im)

            if inpaint_input_mask is not None:
                im = Image.open(str(inpaint_input_mask))
                inpaint_input_mask = np.array(im)

            inpaint_input_image_dict = {
                'image': inpaint_input_image,
                'mask': inpaint_input_mask
            }

        outpaint_selections_arr = []
        for e in outpaint_selections.strip().split(','):
            expansion = e.strip()
            if expansion in outpaint_expansions:
                outpaint_selections_arr.append(expansion)

        image_prompts = []
        image_prompt_config = [
            (cn_img1, cn_stop1, cn_weight1, cn_type1),
            (cn_img2, cn_stop2, cn_weight2, cn_type2),
            (cn_img3, cn_stop3, cn_weight3, cn_type3),
            (cn_img4, cn_stop4, cn_weight4, cn_type4)]
        for config in image_prompt_config:
            cn_img, cn_stop, cn_weight, cn_type = config
            if cn_img is not None:
                im = Image.open(str(cn_img))
                cn_img = np.array(im)
                if cn_stop is None:
                    cn_stop = flags.default_parameters[cn_type][0]
                if cn_weight is None:
                    cn_weight = flags.default_parameters[cn_type][1]
                image_prompts.append((cn_img, cn_stop, cn_weight, cn_type))

        advanced_params = None

        params = ImageGenerationParams(
            prompt=prompt,
            negative_prompt=negative_prompt,
            style_selections=style_selections_arr,
            performance_selection=performance_selection,
            aspect_ratios_selection=aspect_ratios_selection,
            image_number=image_number,
            image_seed=image_seed,
            sharpness=sharpness,
            guidance_scale=guidance_scale,
            base_model_name=base_model_name,
            refiner_model_name=refiner_model_name,
            refiner_switch=refiner_switch,
            loras=loras,
            uov_input_image=uov_input_image,
            uov_method=uov_method,
            upscale_value=uov_upscale_value,
            outpaint_selections=outpaint_selections_arr,
            inpaint_input_image=inpaint_input_image_dict,
            image_prompts=image_prompts,
            advanced_params=advanced_params,
            inpaint_additional_prompt=inpaint_additional_prompt,
            outpaint_distance_left=outpaint_distance_left,
            outpaint_distance_top=outpaint_distance_top,
            outpaint_distance_right=outpaint_distance_right,
            outpaint_distance_bottom=outpaint_distance_bottom,
            save_meta=True,
            meta_scheme='fooocus',
            save_extension='png',
            save_name='',
            require_base64=False,
        )

        print(f"[Predictor Predict] Params: {params.__dict__}")

        async_task = worker_queue.add_task(
            TaskType.text_2_img,
            params)

        if async_task is None:
            print("[Task Queue] The task queue has reached limit")
            raise Exception("The task queue has reached limit.")

        results = blocking_get_task_result(async_task.job_id)

        output_paths: List[Path] = []
        output_seeds: List[str] = []
        for r in results:
            if r.finish_reason == GenerationFinishReason.success and r.im is not None:
                output_seeds.append(r.seed)
                output_paths.append(Path(os.path.join(output_dir, r.im)))

        print(f"[Predictor Predict] Finished with {len(output_paths)} images")

        if len(output_paths) == 0:
            raise Exception("Process failed.")

        return Output(seeds=output_seeds, paths=output_paths)