refactor: remove fooocus api

fooocus_api_version.py

@@ -1 +0,0 @@
-version = '0.4.1.1'

fooocusapi/api.py

@@ -1,41 +0,0 @@
-"""
-Entry for startup fastapi server
-"""
-from fastapi import FastAPI
-from fastapi.staticfiles import StaticFiles
-from fastapi.middleware.cors import CORSMiddleware
-
-import uvicorn
-
-from fooocusapi.utils import file_utils
-from fooocusapi.routes.generate_v1 import secure_router as generate_v1
-from fooocusapi.routes.generate_v2 import secure_router as generate_v2
-from fooocusapi.routes.query import secure_router as query
-from mannequin_to_model import secure_router as mannequin_to_model
-
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allow access from all sources
-    allow_credentials=True,
-    allow_methods=["*"],  # Allow all HTTP methods
-    allow_headers=["*"],  # Allow all request headers
-)
-
-app.mount("/files", StaticFiles(directory=file_utils.output_dir), name="files")
-
-app.include_router(query)
-app.include_router(generate_v1)
-app.include_router(generate_v2)
-app.include_router(mannequin_to_model)
-
-
-def start_app(args):
-    """Start the FastAPI application"""
-    file_utils.STATIC_SERVER_BASE = args.base_url + "/files/"
-    uvicorn.run(
-        app="fooocusapi.api:app",
-        host="0.0.0.0",
-        port=8000,
-        log_level=args.log_level)

fooocusapi/args.py

@@ -1,20 +0,0 @@
-"""
-Do not modify the import order
-"""
-from fooocusapi.base_args import add_base_args
-import ldm_patched.modules.args_parser as args_parser
-
-# Add Fooocus-API args to parser
-add_base_args(args_parser.parser, False)
-
-# Apply Fooocus args
-from args_manager import args_parser
-
-# Override the port default value
-args_parser.parser.set_defaults(
-    port=8888
-)
-
-# Execute args parse again
-args_parser.args = args_parser.parser.parse_args()
-args = args_parser.args

fooocusapi/base_args.py

@@ -1,27 +0,0 @@
-"""
-base_args.py
-"""
-from argparse import ArgumentParser
-
-
-def add_base_args(parser: ArgumentParser, before_prepared: bool):
-    """
-    Add base args for fooocusapi
-    Args:
-        parser: ArgumentParser
-        before_prepared: before prepare environment
-    Returns:
-    """
-    if before_prepared:
-        parser.add_argument("--port", type=int, default=8888, help="Set the listen port, default: 8888")
-
-    parser.add_argument("--host", type=str, default='127.0.0.1', help="Set the listen host, default: 127.0.0.1")
-    parser.add_argument("--base-url", type=str, default=None, help="Set base url for outside visit, default is http://host:port")
-    parser.add_argument("--log-level", type=str, default='info', help="Log info for Uvicorn, default: info")
-    parser.add_argument("--skip-pip", default=False, action="store_true", help="Skip automatic pip install when setup")
-    parser.add_argument("--preload-pipeline", default=False, action="store_true", help="Preload pipeline before start http server")
-    parser.add_argument("--queue-size", type=int, default=100, help="Working queue size, default: 100, generation requests exceeding working queue size will return failure")
-    parser.add_argument("--queue-history", type=int, default=0, help="Finished jobs reserve size, tasks exceeding the limit will be deleted, including output image files, default: 0, means no limit")
-    parser.add_argument('--webhook-url', type=str, default=None, help='The URL to send a POST request when a job is finished')
-    parser.add_argument('--persistent', default=False, action="store_true", help="Store history to db")
-    parser.add_argument("--apikey", type=str, default=None, help="API key for authenticating requests")

fooocusapi/configs/default.py

@@ -1,92 +0,0 @@
-"""
-Static variables for Fooocus API
-"""
-img_generate_responses = {
-    "200": {
-        "description": "PNG bytes if request's 'Accept' header is 'image/png', otherwise JSON",
-        "content": {
-            "application/json": {
-                "example": [{
-                        "base64": "...very long string...",
-                        "seed": "1050625087",
-                        "finish_reason": "SUCCESS",
-                    }]
-            },
-            "application/json async": {
-                "example": {
-                    "job_id": 1,
-                    "job_type": "Text to Image"
-                }
-            },
-            "image/png": {
-                "example": "PNG bytes, what did you expect?"
-            },
-        },
-    }
-}
-
-default_inpaint_engine_version = "v2.6"
-
-default_styles = ["Fooocus V2", "Fooocus Enhance", "Fooocus Sharp"]
-default_base_model_name = "juggernautXL_v8Rundiffusion.safetensors"
-default_refiner_model_name = "None"
-default_refiner_switch = 0.5
-default_loras = [[True, "sd_xl_offset_example-lora_1.0.safetensors", 0.1]]
-default_cfg_scale = 7.0
-default_prompt_negative = ""
-default_aspect_ratio = "1152*896"
-default_sampler = "dpmpp_2m_sde_gpu"
-default_scheduler = "karras"
-
-available_aspect_ratios = [
-    "704*1408",
-    "704*1344",
-    "768*1344",
-    "768*1280",
-    "832*1216",
-    "832*1152",
-    "896*1152",
-    "896*1088",
-    "960*1088",
-    "960*1024",
-    "1024*1024",
-    "1024*960",
-    "1088*960",
-    "1088*896",
-    "1152*896",
-    "1152*832",
-    "1216*832",
-    "1280*768",
-    "1344*768",
-    "1344*704",
-    "1408*704",
-    "1472*704",
-    "1536*640",
-    "1600*640",
-    "1664*576",
-    "1728*576",
-]
-
-uov_methods = [
-    "Disabled",
-    "Vary (Subtle)",
-    "Vary (Strong)",
-    "Upscale (1.5x)",
-    "Upscale (2x)",
-    "Upscale (Fast 2x)",
-    "Upscale (Custom)",
-]
-
-outpaint_expansions = ["Left", "Right", "Top", "Bottom"]
-
-
-def get_aspect_ratio_value(label: str) -> str:
-    """
-    Get aspect ratio
-    Args:
-        label: str, aspect ratio
-
-    Returns:
-
-    """
-    return label.split(" ")[0].replace("×", "*")

fooocusapi/models/common/base.py

@@ -1,189 +0,0 @@
-"""Common models"""
-from typing import List, Tuple
-from enum import Enum
-from fastapi import UploadFile
-from fastapi.exceptions import RequestValidationError
-from pydantic import (
-    ValidationError,
-    ConfigDict,
-    BaseModel,
-    TypeAdapter,
-    Field
-)
-from pydantic_core import InitErrorDetails
-
-from fooocusapi.configs.default import default_loras
-
-
-class PerformanceSelection(str, Enum):
-    """Performance selection"""
-    speed = 'Speed'
-    quality = 'Quality'
-    extreme_speed = 'Extreme Speed'
-    lightning = 'Lightning'
-    hyper_sd = 'Hyper-SD'
-
-
-class Lora(BaseModel):
-    """Common params lora model"""
-    enabled: bool
-    model_name: str
-    weight: float = Field(default=0.5, ge=-2, le=2)
-
-    model_config = ConfigDict(
-        protected_namespaces=('protect_me_', 'also_protect_')
-    )
-
-
-LoraList = TypeAdapter(List[Lora])
-default_loras_model = []
-for lora in default_loras:
-    if lora[0] != 'None':
-        default_loras_model.append(
-            Lora(
-                enabled=lora[0],
-                model_name=lora[1],
-                weight=lora[2])
-        )
-default_loras_json = LoraList.dump_json(default_loras_model)
-
-
-class UpscaleOrVaryMethod(str, Enum):
-    """Upscale or Vary method"""
-    subtle_variation = 'Vary (Subtle)'
-    strong_variation = 'Vary (Strong)'
-    upscale_15 = 'Upscale (1.5x)'
-    upscale_2 = 'Upscale (2x)'
-    upscale_fast = 'Upscale (Fast 2x)'
-    upscale_custom = 'Upscale (Custom)'
-
-
-class OutpaintExpansion(str, Enum):
-    """Outpaint expansion"""
-    left = 'Left'
-    right = 'Right'
-    top = 'Top'
-    bottom = 'Bottom'
-
-
-class ControlNetType(str, Enum):
-    """ControlNet Type"""
-    cn_ip = "ImagePrompt"
-    cn_ip_face = "FaceSwap"
-    cn_canny = "PyraCanny"
-    cn_cpds = "CPDS"
-
-
-class ImagePrompt(BaseModel):
-    """Common params object ImagePrompt"""
-    cn_img: UploadFile | None = Field(default=None)
-    cn_stop: float | None = Field(default=None, ge=0, le=1)
-    cn_weight: float | None = Field(default=None, ge=0, le=2, description="None for default value")
-    cn_type: ControlNetType = Field(default=ControlNetType.cn_ip)
-
-
-class DescribeImageType(str, Enum):
-    """Image type for image to prompt"""
-    photo = 'Photo'
-    anime = 'Anime'
-
-
-class ImageMetaScheme(str, Enum):
-    """Scheme for save image meta
-    Attributes:
-        Fooocus: json format
-        A111: string
-    """
-    Fooocus = 'fooocus'
-    A111 = 'a111'
-
-
-def style_selection_parser(style_selections: str | List[str]) -> List[str]:
-    """
-    Parse style selections, Convert to list
-    Args:
-        style_selections: str, comma separated Fooocus style selections
-        e.g. Fooocus V2, Fooocus Enhance, Fooocus Sharp
-    Returns:
-        List[str]
-    """
-    style_selection_arr: List[str] = []
-    if style_selections is None or len(style_selections) == 0:
-        return []
-    for part in style_selections:
-        if len(part) > 0:
-            for s in part.split(','):
-                style = s.strip()
-                style_selection_arr.append(style)
-    return style_selection_arr
-
-
-def lora_parser(loras: str) -> List[Lora]:
-    """
-    Parse lora config, Convert to list
-    Args:
-        loras: a json string for loras
-    Returns:
-        List[Lora]
-    """
-    loras_model: List[Lora] = []
-    if loras is None or len(loras) == 0:
-        return loras_model
-    try:
-        loras_model = LoraList.validate_json(loras)
-        return loras_model
-    except ValidationError as ve:
-        errs = ve.errors()
-        raise RequestValidationError from errs
-
-
-def outpaint_selections_parser(outpaint_selections: str | list[str]) -> List[OutpaintExpansion]:
-    """
-    Parse outpaint selections, Convert to list
-    Args:
-        outpaint_selections: str, comma separated Left, Right, Top, Bottom
-        e.g. Left, Right, Top, Bottom
-    Returns:
-        List[OutpaintExpansion]
-    """
-    outpaint_selections_arr: List[OutpaintExpansion] = []
-    if outpaint_selections is None or len(outpaint_selections) == 0:
-        return []
-    for part in outpaint_selections:
-        if len(part) > 0:
-            for s in part.split(','):
-                try:
-                    expansion = OutpaintExpansion(s)
-                    outpaint_selections_arr.append(expansion)
-                except ValueError:
-                    errs = InitErrorDetails(
-                        type='enum',
-                        loc=tuple('outpaint_selections'),
-                        input=outpaint_selections,
-                        ctx={
-                            'expected': "str, comma separated Left, Right, Top, Bottom"
-                        })
-                    raise RequestValidationError from errs
-    return outpaint_selections_arr
-
-
-def image_prompt_parser(image_prompts_config: List[Tuple]) -> List[ImagePrompt]:
-    """
-    Image prompt parser, Convert to List[ImagePrompt]
-    Args:
-        image_prompts_config: List[Tuple]
-        e.g. ('image1.jpg', 0.5, 1.0, 'normal'), ('image2.jpg', 0.5, 1.0, 'normal')
-    returns:
-        List[ImagePrompt]
-    """
-    image_prompts: List[ImagePrompt] = []
-    if image_prompts_config is None or len(image_prompts_config) == 0:
-        return []
-    for config in image_prompts_config:
-        cn_img, cn_stop, cn_weight, cn_type = config
-        image_prompts.append(ImagePrompt(
-            cn_img=cn_img,
-            cn_stop=cn_stop,
-            cn_weight=cn_weight,
-            cn_type=cn_type))
-    return image_prompts

fooocusapi/models/common/image_meta.py

@@ -1,118 +0,0 @@
-"""
-Image meta schema
-"""
-from typing import List
-
-from fooocus_version import version
-from pydantic import BaseModel
-
-
-class ImageMeta(BaseModel):
-    """
-    Image meta data model
-    """
-
-    metadata_scheme: str = "fooocus"
-
-    base_model: str
-    base_model_hash: str
-
-    prompt: str
-    full_prompt: List[str]
-    prompt_expansion: str
-
-    negative_prompt: str
-    full_negative_prompt: List[str]
-
-    performance: str
-
-    style: str
-
-    refiner_model: str = "None"
-    refiner_switch: float = 0.5
-
-    loras: List[list]
-
-    resolution: str
-
-    sampler: str = "dpmpp_2m_sde_gpu"
-    scheduler: str = "karras"
-    seed: str
-    adm_guidance: str
-    guidance_scale: float
-    sharpness: float
-    steps: int
-    vae_name: str
-
-    version: str = version
-
-    def __repr__(self):
-        return ""
-
-
-def loras_parser(loras: list) -> list:
-    """
-    Parse lora list
-    """
-    return [
-        [
-            lora[0].rsplit('.', maxsplit=1)[:1][0],
-            lora[1],
-            "hash_not_calculated",
-        ] for lora in loras if lora[0] != 'None' and lora[0] is not None]
-
-
-def image_parse(
-        async_tak: object,
-        task: dict
-) -> dict | str:
-    """
-    Parse image meta data
-    Generate meta data for image from task and async task object
-    Args:
-        async_tak: async task obj
-        task: task obj
-
-    Returns:
-        dict: image meta data
-    """
-    req_param = async_tak.req_param
-    meta = ImageMeta(
-        metadata_scheme=req_param.meta_scheme,
-        base_model=req_param.base_model_name.rsplit('.', maxsplit=1)[:1][0],
-        base_model_hash='',
-        prompt=req_param.prompt,
-        full_prompt=task['positive'],
-        prompt_expansion=task['expansion'],
-        negative_prompt=req_param.negative_prompt,
-        full_negative_prompt=task['negative'],
-        performance=req_param.performance_selection,
-        style=str(req_param.style_selections),
-        refiner_model=req_param.refiner_model_name,
-        refiner_switch=req_param.refiner_switch,
-        loras=loras_parser(req_param.loras),
-        resolution=str(tuple([int(n) for n in req_param.aspect_ratios_selection.split('*')])),
-        sampler=req_param.advanced_params.sampler_name,
-        scheduler=req_param.advanced_params.scheduler_name,
-        seed=str(task['task_seed']),
-        adm_guidance=str((
-            req_param.advanced_params.adm_scaler_positive,
-            req_param.advanced_params.adm_scaler_negative,
-            req_param.advanced_params.adm_scaler_end)),
-        guidance_scale=req_param.guidance_scale,
-        sharpness=req_param.sharpness,
-        steps=-1,
-        vae_name=req_param.advanced_params.vae_name,
-        version=version
-    )
-    if meta.metadata_scheme not in ["fooocus", "a111"]:
-        meta.metadata_scheme = "fooocus"
-    if meta.metadata_scheme == "fooocus":
-        meta_dict = meta.model_dump()
-        for i, lora in enumerate(meta.loras):
-            attr_name = f"lora_combined_{i+1}"
-            lr = [str(x) for x in lora]
-            meta_dict[attr_name] = f"{lr[0]} : {lr[1]}"
-    else:
-        meta_dict = meta.model_dump()
-    return meta_dict

fooocusapi/models/common/requests.py

@@ -1,132 +0,0 @@
-"""Common model for requests"""
-from typing import List
-from pydantic import (
-    BaseModel,
-    Field,
-    ValidationError
-)
-
-from modules.config import (
-    default_sampler,
-    default_scheduler,
-    default_prompt,
-    default_prompt_negative,
-    default_aspect_ratio,
-    default_base_model_name,
-    default_refiner_model_name,
-    default_refiner_switch,
-    default_cfg_scale,
-    default_styles,
-    default_overwrite_step,
-    default_inpaint_engine_version,
-    default_overwrite_switch,
-    default_cfg_tsnr,
-    default_sample_sharpness,
-    default_vae,
-    default_clip_skip
-)
-
-from modules.flags import clip_skip_max
-
-from fooocusapi.models.common.base import (
-    PerformanceSelection,
-    Lora,
-    default_loras_model
-)
-
-default_aspect_ratio = default_aspect_ratio.split(" ")[0].replace("×", "*")
-
-
-class QueryJobRequest(BaseModel):
-    """Query job request"""
-    job_id: str = Field(description="Job ID to query")
-    require_step_preview: bool = Field(
-        default=False,
-        description="Set to true will return preview image of generation steps at current time")
-
-
-class AdvancedParams(BaseModel):
-    """Common params object AdvancedParams"""
-    disable_preview: bool = Field(False, description="Disable preview during generation")
-    disable_intermediate_results: bool = Field(False, description="Disable intermediate results")
-    disable_seed_increment: bool = Field(False, description="Disable Seed Increment")
-    adm_scaler_positive: float = Field(1.5, description="Positive ADM Guidance Scaler", ge=0.1, le=3.0)
-    adm_scaler_negative: float = Field(0.8, description="Negative ADM Guidance Scaler", ge=0.1, le=3.0)
-    adm_scaler_end: float = Field(0.3, description="ADM Guidance End At Step", ge=0.0, le=1.0)
-    adaptive_cfg: float = Field(default_cfg_tsnr, description="CFG Mimicking from TSNR", ge=1.0, le=30.0)
-    clip_skip: int = Field(default_clip_skip, description="Clip Skip", ge=1, le=clip_skip_max)
-    sampler_name: str = Field(default_sampler, description="Sampler")
-    scheduler_name: str = Field(default_scheduler, description="Scheduler")
-    overwrite_step: int = Field(default_overwrite_step, description="Forced Overwrite of Sampling Step", ge=-1, le=200)
-    overwrite_switch: float = Field(default_overwrite_switch, description="Forced Overwrite of Refiner Switch Step", ge=-1, le=1)
-    overwrite_width: int = Field(-1, description="Forced Overwrite of Generating Width", ge=-1, le=2048)
-    overwrite_height: int = Field(-1, description="Forced Overwrite of Generating Height", ge=-1, le=2048)
-    overwrite_vary_strength: float = Field(-1, description='Forced Overwrite of Denoising Strength of "Vary"', ge=-1, le=1.0)
-    overwrite_upscale_strength: float = Field(-1, description='Forced Overwrite of Denoising Strength of "Upscale"', ge=-1, le=1.0)
-    mixing_image_prompt_and_vary_upscale: bool = Field(False, description="Mixing Image Prompt and Vary/Upscale")
-    mixing_image_prompt_and_inpaint: bool = Field(False, description="Mixing Image Prompt and Inpaint")
-    debugging_cn_preprocessor: bool = Field(False, description="Debug Preprocessors")
-    skipping_cn_preprocessor: bool = Field(False, description="Skip Preprocessors")
-    canny_low_threshold: int = Field(64, description="Canny Low Threshold", ge=1, le=255)
-    canny_high_threshold: int = Field(128, description="Canny High Threshold", ge=1, le=255)
-    refiner_swap_method: str = Field('joint', description="Refiner swap method")
-    controlnet_softness: float = Field(0.25, description="Softness of ControlNet", ge=0.0, le=1.0)
-    freeu_enabled: bool = Field(False, description="FreeU enabled")
-    freeu_b1: float = Field(1.01, description="FreeU B1")
-    freeu_b2: float = Field(1.02, description="FreeU B2")
-    freeu_s1: float = Field(0.99, description="FreeU B3")
-    freeu_s2: float = Field(0.95, description="FreeU B4")
-    debugging_inpaint_preprocessor: bool = Field(False, description="Debug Inpaint Preprocessing")
-    inpaint_disable_initial_latent: bool = Field(False, description="Disable initial latent in inpaint")
-    inpaint_engine: str = Field(default_inpaint_engine_version, description="Inpaint Engine")
-    inpaint_strength: float = Field(1.0, description="Inpaint Denoising Strength", ge=0.0, le=1.0)
-    inpaint_respective_field: float = Field(1.0, description="Inpaint Respective Field", ge=0.0, le=1.0)
-    inpaint_mask_upload_checkbox: bool = Field(False, description="Upload Mask")
-    invert_mask_checkbox: bool = Field(False, description="Invert Mask")
-    inpaint_erode_or_dilate: int = Field(0, description="Mask Erode or Dilate", ge=-64, le=64)
-    black_out_nsfw: bool = Field(False, description="Block out NSFW")
-    vae_name: str = Field(default_vae, description="VAE name")
-
-
-class CommonRequest(BaseModel):
-    """All generate request based on this model"""
-    prompt: str = default_prompt
-    negative_prompt: str = default_prompt_negative
-    style_selections: List[str] = default_styles
-    performance_selection: PerformanceSelection = PerformanceSelection.speed
-    aspect_ratios_selection: str = default_aspect_ratio
-    image_number: int = Field(default=1, description="Image number", ge=1, le=32)
-    image_seed: int = Field(default=-1, description="Seed to generate image, -1 for random")
-    sharpness: float = Field(default=default_sample_sharpness, ge=0.0, le=30.0)
-    guidance_scale: float = Field(default=default_cfg_scale, ge=1.0, le=30.0)
-    base_model_name: str = default_base_model_name
-    refiner_model_name: str = default_refiner_model_name
-    refiner_switch: float = Field(default=default_refiner_switch, description="Refiner Switch At", ge=0.1, le=1.0)
-    loras: List[Lora] = Field(default=default_loras_model)
-    advanced_params: AdvancedParams = AdvancedParams()
-    save_meta: bool = Field(default=True, description="Save meta data")
-    meta_scheme: str = Field(default='fooocus', description="Meta data scheme, one of [fooocus, a111]")
-    save_extension: str = Field(default='png', description="Save extension, one of [png, jpg, webp]")
-    save_name: str = Field(default='', description="Image name for output image, default is job id + seq")
-    read_wildcards_in_order: bool = Field(default=False, description="Read wildcards in order")
-    require_base64: bool = Field(default=False, description="Return base64 data of generated image")
-    async_process: bool = Field(default=False, description="Set to true will run async and return job info for retrieve generation result later")
-    webhook_url: str | None = Field(default='', description="Optional URL for a webhook callback. If provided, the system will send a POST request to this URL upon task completion or failure."
-                                                            " This allows for asynchronous notification of task status.")
-
-
-def advanced_params_parser(advanced_params: str | None) -> AdvancedParams:
-    """
-    Parse advanced params, Convert to AdvancedParams
-    Args:
-        advanced_params: str, json format
-    Returns:
-        AdvancedParams object, if validate error return default value
-    """
-    if advanced_params is not None and len(advanced_params) > 0:
-        try:
-            advanced_params_obj = AdvancedParams.__pydantic_validator__.validate_json(advanced_params)
-            return AdvancedParams(**advanced_params_obj)
-        except ValidationError:
-            return AdvancedParams()
-    return AdvancedParams()

fooocusapi/models/common/response.py

@@ -1,90 +0,0 @@
-"""Fooocus API models for response"""
-from typing import List
-
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field
-)
-
-from fooocusapi.models.common.task import (
-    GeneratedImageResult,
-    AsyncJobStage
-)
-from fooocusapi.task_queue import TaskType
-
-
-class DescribeImageResponse(BaseModel):
-    """
-    describe image response
-    """
-    describe: str
-
-
-class AsyncJobResponse(BaseModel):
-    """
-    Async job response
-    Attributes:
-        job_id: Job ID
-        job_type: Job type
-        job_stage: Job stage
-        job_progress: Job progress, 0-100
-        job_status: Job status
-        job_step_preview: Job step preview
-        job_result: Job result
-    """
-    job_id: str = Field(description="Job ID")
-    job_type: TaskType = Field(description="Job type")
-    job_stage: AsyncJobStage = Field(description="Job running stage")
-    job_progress: int = Field(description="Job running progress, 100 is for finished.")
-    job_status: str | None = Field(None, description="Job running status in text")
-    job_step_preview: str | None = Field(None, description="Preview image of generation steps at current time, as base64 image")
-    job_result: List[GeneratedImageResult] | None = Field(None, description="Job generation result")
-
-
-class JobQueueInfo(BaseModel):
-    """
-    job queue info
-    Attributes:
-        running_size: int, The current running and waiting job count
-        finished_size: int, The current finished job count
-        last_job_id: str, Last submit generation job id
-    """
-    running_size: int = Field(description="The current running and waiting job count")
-    finished_size: int = Field(description="Finished job count (after auto clean)")
-    last_job_id: str | None = Field(description="Last submit generation job id")
-
-
-# TODO May need more detail fields, will add later when someone need
-class JobHistoryInfo(BaseModel):
-    """
-    job history info
-    """
-    job_id: str
-    is_finished: bool = False
-
-
-# Response model for the historical tasks
-class JobHistoryResponse(BaseModel):
-    """
-    job history response
-    """
-    queue: List[JobHistoryInfo] = []
-    history: List[JobHistoryInfo] = []
-
-
-class AllModelNamesResponse(BaseModel):
-    """
-    all model list response
-    """
-    model_filenames: List[str] = Field(description="All available model filenames")
-    lora_filenames: List[str] = Field(description="All available lora filenames")
-
-    model_config = ConfigDict(
-        protected_namespaces=('protect_me_', 'also_protect_')
-    )
-
-
-class StopResponse(BaseModel):
-    """stop task response"""
-    msg: str

fooocusapi/models/common/task.py

@@ -1,60 +0,0 @@
-"""
-Task and job related models
-"""
-from enum import Enum
-from pydantic import (
-    BaseModel,
-    Field
-)
-
-
-class TaskType(str, Enum):
-    """
-    Task type object
-    """
-    text_2_img = 'Text to Image'
-    img_uov = 'Image Upscale or Variation'
-    img_inpaint_outpaint = 'Image Inpaint or Outpaint'
-    img_prompt = 'Image Prompt'
-    not_found = 'Not Found'
-
-
-class GenerationFinishReason(str, Enum):
-    """
-    Generation finish reason
-    """
-    success = 'SUCCESS'
-    queue_is_full = 'QUEUE_IS_FULL'
-    user_cancel = 'USER_CANCEL'
-    error = 'ERROR'
-
-
-class ImageGenerationResult:
-    """
-    Image generation result
-    """
-    def __init__(self, im: str | None, seed: str, finish_reason: GenerationFinishReason):
-        self.im = im
-        self.seed = seed
-        self.finish_reason = finish_reason
-
-
-class AsyncJobStage(str, Enum):
-    """
-    Async job stage
-    """
-    waiting = 'WAITING'
-    running = 'RUNNING'
-    success = 'SUCCESS'
-    error = 'ERROR'
-
-
-class GeneratedImageResult(BaseModel):
-    """
-    Generated images result
-    """
-    base64: str | None = Field(
-        description="Image encoded in base64, or null if finishReason is not 'SUCCESS', only return when request require base64")
-    url: str | None = Field(description="Image file static serve url, or null if finishReason is not 'SUCCESS'")
-    seed: str = Field(description="The seed associated with this image")
-    finish_reason: GenerationFinishReason

fooocusapi/models/requests_v1.py

@@ -1,274 +0,0 @@
-"""
-requests models for v1 endpoints
-"""
-from typing import List
-from fastapi.params import File
-from fastapi import (
-    UploadFile,
-    Form
-)
-from fooocusapi.models.common.requests import (
-    CommonRequest,
-    advanced_params_parser
-)
-from fooocusapi.models.common.base import (
-    ImagePrompt,
-    ControlNetType,
-    OutpaintExpansion,
-    UpscaleOrVaryMethod,
-    PerformanceSelection
-)
-
-from fooocusapi.models.common.base import (
-    style_selection_parser,
-    lora_parser,
-    outpaint_selections_parser,
-    image_prompt_parser,
-    default_loras_json
-)
-
-from fooocusapi.configs.default import (
-    default_prompt_negative,
-    default_aspect_ratio,
-    default_base_model_name,
-    default_refiner_model_name,
-    default_refiner_switch,
-    default_cfg_scale,
-    default_styles,
-)
-
-
-class ImgUpscaleOrVaryRequest(CommonRequest):
-    """
-    Request for image upscale or variation
-    Attributes:
-        input_image: Input image
-        uov_method: Upscale or variation method
-        upscale_value: upscale value
-    Functions:
-        as_form: Convert request to form data
-    """
-    input_image: UploadFile
-    uov_method: UpscaleOrVaryMethod
-    upscale_value: float | None
-
-    @classmethod
-    def as_form(
-            cls,
-            input_image: UploadFile = Form(description="Init image for upscale or outpaint"),
-            uov_method: UpscaleOrVaryMethod = Form(),
-            upscale_value: float | None = Form(None, description="Upscale custom value, None for default value", ge=1.0, le=5.0),
-            prompt: str = Form(''),
-            negative_prompt: str = Form(default_prompt_negative),
-            style_selections: List[str] = Form(default_styles, description="Fooocus style selections, separated by comma"),
-            performance_selection: PerformanceSelection = Form(PerformanceSelection.speed, description="Performance Selection, one of 'Speed','Quality','Extreme Speed'"),
-            aspect_ratios_selection: str = Form(default_aspect_ratio, description="Aspect Ratios Selection, default 1152*896"),
-            image_number: int = Form(default=1, description="Image number", ge=1, le=32),
-            image_seed: int = Form(default=-1, description="Seed to generate image, -1 for random"),
-            sharpness: float = Form(default=2.0, ge=0.0, le=30.0),
-            guidance_scale: float = Form(default=default_cfg_scale, ge=1.0, le=30.0),
-            base_model_name: str = Form(default_base_model_name, description="checkpoint file name"),
-            refiner_model_name: str = Form(default_refiner_model_name, description="refiner file name"),
-            refiner_switch: float = Form(default=default_refiner_switch, description="Refiner Switch At", ge=0.1, le=1.0),
-            loras: str | None = Form(default=default_loras_json, description='Lora config in JSON. Format as [{"model_name": "sd_xl_offset_example-lora_1.0.safetensors", "weight": 0.5}]'),
-            advanced_params: str | None = Form(default=None, description="Advanced parameters in JSON"),
-            save_meta: bool = Form(default=False, description="Save metadata to image"),
-            meta_scheme: str = Form(default='fooocus', description="Metadata scheme, one of 'fooocus', 'a111'"),
-            save_extension: str = Form(default="png", description="Save extension, png, jpg or webp"),
-            save_name: str = Form(default="", description="Save name, empty for auto generate"),
-            require_base64: bool = Form(default=False, description="Return base64 data of generated image"),
-            read_wildcards_in_order: bool = Form(default=False, description="Read wildcards in order"),
-            async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generation result later"),
-            webhook_url: str = Form(default="", description="Webhook url for generation result"),
-    ):
-        style_selection_arr = style_selection_parser(style_selections)
-        loras_model = lora_parser(loras)
-        advanced_params_obj = advanced_params_parser(advanced_params)
-
-        return cls(
-            input_image=input_image, uov_method=uov_method, upscale_value=upscale_value,
-            prompt=prompt, negative_prompt=negative_prompt, style_selections=style_selection_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_model, advanced_params=advanced_params_obj, save_meta=save_meta, meta_scheme=meta_scheme,
-            save_extension=save_extension, save_name=save_name, require_base64=require_base64,
-            read_wildcards_in_order=read_wildcards_in_order, async_process=async_process, webhook_url=webhook_url)
-
-
-class ImgInpaintOrOutpaintRequest(CommonRequest):
-    """
-    Image Inpaint or Outpaint Request
-    """
-    input_image: UploadFile | None
-    input_mask: UploadFile | None
-    inpaint_additional_prompt: str | None
-    outpaint_selections: List[OutpaintExpansion]
-    outpaint_distance_left: int
-    outpaint_distance_right: int
-    outpaint_distance_top: int
-    outpaint_distance_bottom: int
-
-    @classmethod
-    def as_form(
-            cls,
-            input_image: UploadFile = Form(description="Init image for inpaint or outpaint"),
-            input_mask: UploadFile = Form(File(None), description="Inpaint or outpaint mask"),
-            inpaint_additional_prompt: str | None = Form("", description="Describe what you want to inpaint"),
-            outpaint_selections: List[str] = Form([], description="Outpaint expansion selections, literal 'Left', 'Right', 'Top', 'Bottom' separated by comma"),
-            outpaint_distance_left: int = Form(default=0, description="Set outpaint left distance, -1 for default"),
-            outpaint_distance_right: int = Form(default=0, description="Set outpaint right distance, -1 for default"),
-            outpaint_distance_top: int = Form(default=0, description="Set outpaint top distance, -1 for default"),
-            outpaint_distance_bottom: int = Form(default=0, description="Set outpaint bottom distance, -1 for default"),
-            prompt: str = Form(''),
-            negative_prompt: str = Form(default_prompt_negative),
-            style_selections: List[str] = Form(default_styles, description="Fooocus style selections, separated by comma"),
-            performance_selection: PerformanceSelection = Form(PerformanceSelection.speed, description="Performance Selection, one of 'Speed','Quality','Extreme Speed'"),
-            aspect_ratios_selection: str = Form(default_aspect_ratio, description="Aspect Ratios Selection, default 1152*896"),
-            image_number: int = Form(default=1, description="Image number", ge=1, le=32),
-            image_seed: int = Form(default=-1, description="Seed to generate image, -1 for random"),
-            sharpness: float = Form(default=2.0, ge=0.0, le=30.0),
-            guidance_scale: float = Form(default=default_cfg_scale, ge=1.0, le=30.0),
-            base_model_name: str = Form(default_base_model_name),
-            refiner_model_name: str = Form(default_refiner_model_name),
-            refiner_switch: float = Form(default=default_refiner_switch, description="Refiner Switch At", ge=0.1, le=1.0),
-            loras: str | None = Form(default=default_loras_json, description='Lora config in JSON. Format as [{"model_name": "sd_xl_offset_example-lora_1.0.safetensors", "weight": 0.5}]'),
-            advanced_params: str | None = Form(default=None, description="Advanced parameters in JSON"),
-            save_meta: bool = Form(default=False, description="Save metadata to image"),
-            meta_scheme: str = Form(default='fooocus', description="Metadata scheme, one of 'fooocus', 'a111'"),
-            save_extension: str = Form(default="png", description="Save extension, png, jpg or webp"),
-            save_name: str = Form(default="", description="Save name, empty for auto generate"),
-            require_base64: bool = Form(default=False, description="Return base64 data of generated image"),
-            read_wildcards_in_order: bool = Form(default=False, description="Read wildcards in order"),
-            async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generation result later"),
-            webhook_url: str = Form(default="", description="Webhook url for generation result"),
-    ):
-        if isinstance(input_mask, File):
-            input_mask = None
-
-        outpaint_selections_arr = outpaint_selections_parser(outpaint_selections)
-        style_selection_arr = style_selection_parser(style_selections)
-        loras_model = lora_parser(loras)
-        advanced_params_obj = advanced_params_parser(advanced_params)
-
-        return cls(
-            input_image=input_image, input_mask=input_mask, inpaint_additional_prompt=inpaint_additional_prompt,
-            outpaint_selections=outpaint_selections_arr, outpaint_distance_left=outpaint_distance_left,
-            outpaint_distance_right=outpaint_distance_right, outpaint_distance_top=outpaint_distance_top,
-            outpaint_distance_bottom=outpaint_distance_bottom, prompt=prompt, negative_prompt=negative_prompt, style_selections=style_selection_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_model, advanced_params=advanced_params_obj, save_meta=save_meta, meta_scheme=meta_scheme,
-            save_extension=save_extension, save_name=save_name, require_base64=require_base64,
-            read_wildcards_in_order=read_wildcards_in_order, async_process=async_process, webhook_url=webhook_url)
-
-
-class ImgPromptRequest(ImgInpaintOrOutpaintRequest):
-    """
-    Image Prompt Request
-    """
-    image_prompts: List[ImagePrompt]
-
-    @classmethod
-    def as_form(
-            cls,
-            input_image: UploadFile = Form(File(None), description="Init image for inpaint or outpaint"),
-            input_mask: UploadFile = Form(File(None), description="Inpaint or outpaint mask"),
-            inpaint_additional_prompt: str | None = Form(None, description="Describe what you want to inpaint"),
-            outpaint_selections: List[str] = Form([], description="Outpaint expansion selections, literal 'Left', 'Right', 'Top', 'Bottom' separated by comma"),
-            outpaint_distance_left: int = Form(default=0, description="Set outpaint left distance, 0 for default"),
-            outpaint_distance_right: int = Form(default=0, description="Set outpaint right distance, 0 for default"),
-            outpaint_distance_top: int = Form(default=0, description="Set outpaint top distance, 0 for default"),
-            outpaint_distance_bottom: int = Form(default=0, description="Set outpaint bottom distance, 0 for default"),
-            cn_img1: UploadFile = Form(File(None), description="Input image for image prompt"),
-            cn_stop1: float | None = Form(
-                default=None, ge=0, le=1, description="Stop at for image prompt, None for default value"),
-            cn_weight1: float | None = Form(
-                default=None, ge=0, le=2, description="Weight for image prompt, None for default value"),
-            cn_type1: ControlNetType = Form(
-                default=ControlNetType.cn_ip, description="ControlNet type for image prompt"),
-            cn_img2: UploadFile = Form(
-                File(None), description="Input image for image prompt"),
-            cn_stop2: float | None = Form(
-                default=None, ge=0, le=1, description="Stop at for image prompt, None for default value"),
-            cn_weight2: float | None = Form(
-                default=None, ge=0, le=2, description="Weight for image prompt, None for default value"),
-            cn_type2: ControlNetType = Form(
-                default=ControlNetType.cn_ip, description="ControlNet type for image prompt"),
-            cn_img3: UploadFile = Form(
-                File(None), description="Input image for image prompt"),
-            cn_stop3: float | None = Form(
-                default=None, ge=0, le=1, description="Stop at for image prompt, None for default value"),
-            cn_weight3: float | None = Form(
-                default=None, ge=0, le=2, description="Weight for image prompt, None for default value"),
-            cn_type3: ControlNetType = Form(
-                default=ControlNetType.cn_ip, description="ControlNet type for image prompt"),
-            cn_img4: UploadFile = Form(
-                File(None), description="Input image for image prompt"),
-            cn_stop4: float | None = Form(
-                default=None, ge=0, le=1, description="Stop at for image prompt, None for default value"),
-            cn_weight4: float | None = Form(
-                default=None, ge=0, le=2, description="Weight for image prompt, None for default value"),
-            cn_type4: ControlNetType = Form(
-                default=ControlNetType.cn_ip, description="ControlNet type for image prompt"),
-            prompt: str = Form(''),
-            negative_prompt: str = Form(default_prompt_negative),
-            style_selections: List[str] = Form(default_styles, description="Fooocus style selections, separated by comma"),
-            performance_selection: PerformanceSelection = Form(
-                PerformanceSelection.speed),
-            aspect_ratios_selection: str = Form(default_aspect_ratio),
-            image_number: int = Form(
-                default=1, description="Image number", ge=1, le=32),
-            image_seed: int = Form(default=-1, description="Seed to generate image, -1 for random"),
-            sharpness: float = Form(default=2.0, ge=0.0, le=30.0),
-            guidance_scale: float = Form(default=default_cfg_scale, ge=1.0, le=30.0),
-            base_model_name: str = Form(default_base_model_name),
-            refiner_model_name: str = Form(default_refiner_model_name),
-            refiner_switch: float = Form(default=default_refiner_switch, description="Refiner Switch At", ge=0.1, le=1.0),
-            loras: str | None = Form(default=default_loras_json, description='Lora config in JSON. Format as [{"model_name": "sd_xl_offset_example-lora_1.0.safetensors", "weight": 0.5}]'),
-            advanced_params: str | None = Form(default=None, description="Advanced parameters in JSON"),
-            save_meta: bool = Form(default=False, description="Save metadata to image"),
-            meta_scheme: str = Form(default='fooocus', description="Metadata scheme, one of 'fooocus', 'a111'"),
-            save_extension: str = Form(default="png", description="Save extension, png, jpg or webp"),
-            save_name: str = Form(default="", description="Save name, empty for auto generate"),
-            require_base64: bool = Form(default=False, description="Return base64 data of generated image"),
-            read_wildcards_in_order: bool = Form(default=False, description="Read wildcards in order"),
-            async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generation result later"),
-            webhook_url: str = Form(default="", description="Webhook url for generation result"),
-    ):
-        if isinstance(input_image, File):
-            input_image = None
-        if isinstance(input_mask, File):
-            input_mask = None
-        if isinstance(cn_img1, File):
-            cn_img1 = None
-        if isinstance(cn_img2, File):
-            cn_img2 = None
-        if isinstance(cn_img3, File):
-            cn_img3 = None
-        if isinstance(cn_img4, File):
-            cn_img4 = None
-
-        outpaint_selections_arr = outpaint_selections_parser(outpaint_selections)
-
-        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)]
-        image_prompts = image_prompt_parser(image_prompt_config)
-        style_selection_arr = style_selection_parser(style_selections)
-        loras_model = lora_parser(loras)
-        advanced_params_obj = advanced_params_parser(advanced_params)
-
-        return cls(
-            input_image=input_image, input_mask=input_mask, inpaint_additional_prompt=inpaint_additional_prompt, outpaint_selections=outpaint_selections_arr,
-            outpaint_distance_left=outpaint_distance_left, outpaint_distance_right=outpaint_distance_right, outpaint_distance_top=outpaint_distance_top, outpaint_distance_bottom=outpaint_distance_bottom,
-            image_prompts=image_prompts, prompt=prompt, negative_prompt=negative_prompt, style_selections=style_selection_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_model, advanced_params=advanced_params_obj, save_meta=save_meta, meta_scheme=meta_scheme,
-            save_extension=save_extension, save_name=save_name, require_base64=require_base64,
-            read_wildcards_in_order=read_wildcards_in_order, async_process=async_process, webhook_url=webhook_url)

fooocusapi/models/requests_v2.py

@@ -1,50 +0,0 @@
-"""V2 API models"""
-from typing import List
-from pydantic import BaseModel, Field
-from fooocusapi.models.common.requests import CommonRequest
-from fooocusapi.models.common.base import (
-    ControlNetType,
-    OutpaintExpansion,
-    ImagePrompt,
-    UpscaleOrVaryMethod
-)
-
-
-class ImagePromptJson(BaseModel):
-    """Image prompt for V2 API"""
-    cn_img: str | None = Field(None, description="Input image for image prompt as base64")
-    cn_stop: float | None = Field(0, ge=0, le=1, description="Stop at for image prompt, 0 for default value")
-    cn_weight: float | None = Field(0, ge=0, le=2, description="Weight for image prompt, 0 for default value")
-    cn_type: ControlNetType = Field(default=ControlNetType.cn_ip, description="ControlNet type for image prompt")
-
-
-class ImgInpaintOrOutpaintRequestJson(CommonRequest):
-    """image inpaint or outpaint request"""
-    input_image: str = Field('', description="Init image for inpaint or outpaint as base64")
-    input_mask: str | None = Field('', description="Inpaint or outpaint mask as base64")
-    inpaint_additional_prompt: str | None = Field('', description="Describe what you want to inpaint")
-    outpaint_selections: List[OutpaintExpansion] = []
-    outpaint_distance_left: int | None = Field(-1, description="Set outpaint left distance")
-    outpaint_distance_right: int | None = Field(-1, description="Set outpaint right distance")
-    outpaint_distance_top: int | None = Field(-1, description="Set outpaint top distance")
-    outpaint_distance_bottom: int | None = Field(-1, description="Set outpaint bottom distance")
-    image_prompts: List[ImagePromptJson | ImagePrompt] = []
-
-
-class ImgPromptRequestJson(ImgInpaintOrOutpaintRequestJson):
-    """img prompt request json"""
-    input_image: str | None = Field(None, description="Init image for inpaint or outpaint as base64")
-    image_prompts: List[ImagePromptJson | ImagePrompt]
-
-
-class Text2ImgRequestWithPrompt(CommonRequest):
-    """text to image request with prompt"""
-    image_prompts: List[ImagePromptJson] = []
-
-
-class ImgUpscaleOrVaryRequestJson(CommonRequest):
-    """img upscale or vary request json"""
-    uov_method: UpscaleOrVaryMethod = UpscaleOrVaryMethod.upscale_2
-    upscale_value: float | None = Field(1.0, ge=1.0, le=5.0, description="Upscale custom value, 1.0 for default value")
-    input_image: str = Field(description="Init image for upscale or outpaint as base64")
-    image_prompts: List[ImagePromptJson | ImagePrompt] = []

fooocusapi/parameters.py

@@ -1,94 +0,0 @@
-from typing import Dict, List, Tuple
-import numpy as np
-import copy
-
-from fooocusapi.models.common.requests import AdvancedParams
-
-
-class ImageGenerationParams:
-    def __init__(
-        self,
-        prompt: str,
-        negative_prompt: str,
-        style_selections: List[str],
-        performance_selection: str,
-        aspect_ratios_selection: str,
-        image_number: int,
-        image_seed: int | None,
-        sharpness: float,
-        guidance_scale: float,
-        base_model_name: str,
-        refiner_model_name: str,
-        refiner_switch: float,
-        loras: List[Tuple[str, float]],
-        uov_input_image: np.ndarray | None,
-        uov_method: str,
-        upscale_value: float | None,
-        outpaint_selections: List[str],
-        outpaint_distance_left: int,
-        outpaint_distance_right: int,
-        outpaint_distance_top: int,
-        outpaint_distance_bottom: int,
-        inpaint_input_image: Dict[str, np.ndarray] | None,
-        inpaint_additional_prompt: str | None,
-        image_prompts: List[Tuple[np.ndarray, float, float, str]],
-        advanced_params: List[any] | None,
-        save_extension: str,
-        save_meta: bool,
-        meta_scheme: str,
-        save_name: str,
-        require_base64: bool,
-    ):
-        self.prompt = prompt
-        self.negative_prompt = negative_prompt
-        self.style_selections = style_selections
-        self.performance_selection = performance_selection
-        self.aspect_ratios_selection = aspect_ratios_selection
-        self.image_number = image_number
-        self.image_seed = image_seed
-        self.sharpness = sharpness
-        self.guidance_scale = guidance_scale
-        self.base_model_name = base_model_name
-        self.refiner_model_name = refiner_model_name
-        self.refiner_switch = refiner_switch
-        self.loras = loras
-        self.uov_input_image = uov_input_image
-        self.uov_method = uov_method
-        self.upscale_value = upscale_value
-        self.outpaint_selections = outpaint_selections
-        self.outpaint_distance_left = outpaint_distance_left
-        self.outpaint_distance_right = outpaint_distance_right
-        self.outpaint_distance_top = outpaint_distance_top
-        self.outpaint_distance_bottom = outpaint_distance_bottom
-        self.inpaint_input_image = inpaint_input_image
-        self.inpaint_additional_prompt = inpaint_additional_prompt
-        self.image_prompts = image_prompts
-        self.save_extension = save_extension
-        self.save_meta = save_meta
-        self.meta_scheme = meta_scheme
-        self.save_name = save_name
-        self.require_base64 = require_base64
-        self.advanced_params = advanced_params
-
-        if self.advanced_params is None:
-            self.advanced_params = AdvancedParams()
-
-            # Auto set mixing_image_prompt_and_inpaint to True
-            if len(self.image_prompts) > 0 and self.inpaint_input_image is not None:
-                print("Mixing Image Prompts and Inpaint Enabled")
-                self.advanced_params.mixing_image_prompt_and_inpaint = True
-            if len(self.image_prompts) > 0 and self.uov_input_image is not None:
-                print("Mixing Image Prompts and Vary Upscale Enabled")
-                self.advanced_params.mixing_image_prompt_and_vary_upscale = True
-
-    def to_dict(self):
-        """
-        Convert the ImageGenerationParams object to a dictionary.
-        Args:
-            self:
-
-        Returns:
-            self to dict
-        """
-        obj_dict = copy.deepcopy(self)
-        return obj_dict.__dict__

fooocusapi/routes/generate_v1.py

@@ -1,186 +0,0 @@
-"""Generate API V1 routes
-
-"""
-from typing import List, Optional
-from fastapi import APIRouter, Depends, Header, Query, UploadFile
-from fastapi.params import File
-
-from modules.util import HWC3
-
-from fooocusapi.models.common.base import DescribeImageType
-from fooocusapi.utils.api_utils import api_key_auth
-
-from fooocusapi.models.common.requests import CommonRequest as Text2ImgRequest
-from fooocusapi.models.requests_v1 import (
-    ImgUpscaleOrVaryRequest,
-    ImgPromptRequest,
-    ImgInpaintOrOutpaintRequest
-)
-from fooocusapi.models.common.response import (
-    AsyncJobResponse,
-    GeneratedImageResult,
-    DescribeImageResponse,
-    StopResponse
-)
-from fooocusapi.utils.call_worker import call_worker
-from fooocusapi.utils.img_utils import read_input_image
-from fooocusapi.configs.default import img_generate_responses
-from fooocusapi.worker import process_stop
-
-
-secure_router = APIRouter(
-    dependencies=[Depends(api_key_auth)]
-)
-
-
-def stop_worker():
-    """Interrupt worker process"""
-    process_stop()
-
-
-@secure_router.post(
-        path="/v1/generation/text-to-image",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
-def text2img_generation(
-    req: Text2ImgRequest,
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nText to Image Generation\n
-    A text to image generation endpoint
-    Arguments:
-        req {Text2ImgRequest} -- Text to image generation request
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v1/generation/image-upscale-vary",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
-def img_upscale_or_vary(
-    input_image: UploadFile,
-    req: ImgUpscaleOrVaryRequest = Depends(ImgUpscaleOrVaryRequest.as_form),
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nImage upscale or vary\n
-    Image upscale or vary
-    Arguments:
-        input_image {UploadFile} -- Input image file
-        req {ImgUpscaleOrVaryRequest} -- Request body
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v1/generation/image-inpaint-outpaint",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
-def img_inpaint_or_outpaint(
-    input_image: UploadFile,
-    req: ImgInpaintOrOutpaintRequest = Depends(ImgInpaintOrOutpaintRequest.as_form),
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nInpaint or outpaint\n
-    Inpaint or outpaint
-    Arguments:
-        input_image {UploadFile} -- Input image file
-        req {ImgInpaintOrOutpaintRequest} -- Request body
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v1/generation/image-prompt",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV1"])
-def img_prompt(
-    cn_img1: Optional[UploadFile] = File(None),
-    req: ImgPromptRequest = Depends(ImgPromptRequest.as_form),
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nImage Prompt\n
-    Image Prompt
-    A prompt-based image generation.
-    Arguments:
-        cn_img1 {UploadFile} -- Input image file
-        req {ImgPromptRequest} -- Request body
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v1/tools/describe-image",
-        response_model=DescribeImageResponse,
-        tags=["GenerateV1"])
-def describe_image(
-    image: UploadFile,
-    image_type: DescribeImageType = Query(
-        DescribeImageType.photo,
-        description="Image type, 'Photo' or 'Anime'")):
-    """\nDescribe image\n
-    Describe image, Get tags from an image
-    Arguments:
-        image {UploadFile} -- Image to get tags
-        image_type {DescribeImageType} -- Image type, 'Photo' or 'Anime'
-    Returns:
-        DescribeImageResponse -- Describe image response, a string
-    """
-    if image_type == DescribeImageType.photo:
-        from extras.interrogate import default_interrogator as default_interrogator_photo
-        interrogator = default_interrogator_photo
-    else:
-        from extras.wd14tagger import default_interrogator as default_interrogator_anime
-        interrogator = default_interrogator_anime
-    img = HWC3(read_input_image(image))
-    result = interrogator(img)
-    return DescribeImageResponse(describe=result)
-
-
-@secure_router.post(
-        path="/v1/generation/stop",
-        response_model=StopResponse,
-        description="Job stopping",
-        tags=["Default"])
-def stop():
-    """Interrupt worker"""
-    stop_worker()
-    return StopResponse(msg="success")

fooocusapi/routes/generate_v2.py

@@ -1,199 +0,0 @@
-"""Generate API V2 routes
-
-"""
-from typing import List
-from fastapi import APIRouter, Depends, Header, Query
-
-from fooocusapi.utils.api_utils import api_key_auth
-from fooocusapi.models.requests_v1 import ImagePrompt
-from fooocusapi.models.requests_v2 import (
-    ImgInpaintOrOutpaintRequestJson,
-    ImgPromptRequestJson,
-    Text2ImgRequestWithPrompt,
-    ImgUpscaleOrVaryRequestJson
-)
-from fooocusapi.models.common.response import (
-    AsyncJobResponse,
-    GeneratedImageResult
-)
-from fooocusapi.utils.call_worker import call_worker
-from fooocusapi.utils.img_utils import base64_to_stream
-from fooocusapi.configs.default import img_generate_responses
-
-
-secure_router = APIRouter(
-    dependencies=[Depends(api_key_auth)]
-)
-
-
-@secure_router.post(
-        path="/v2/generation/text-to-image-with-ip",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV2"])
-def text_to_img_with_ip(
-    req: Text2ImgRequestWithPrompt,
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        default=None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nText to image with prompt\n
-    Text to image with prompt
-    Arguments:
-        req {Text2ImgRequestWithPrompt} -- Text to image generation request
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    default_image_prompt = ImagePrompt(cn_img=None)
-    image_prompts_files: List[ImagePrompt] = []
-    for image_prompt in req.image_prompts:
-        image_prompt.cn_img = base64_to_stream(image_prompt.cn_img)
-        image = ImagePrompt(
-            cn_img=image_prompt.cn_img,
-            cn_stop=image_prompt.cn_stop,
-            cn_weight=image_prompt.cn_weight,
-            cn_type=image_prompt.cn_type)
-        image_prompts_files.append(image)
-
-    while len(image_prompts_files) <= 4:
-        image_prompts_files.append(default_image_prompt)
-
-    req.image_prompts = image_prompts_files
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v2/generation/image-upscale-vary",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV2"])
-def img_upscale_or_vary(
-    req: ImgUpscaleOrVaryRequestJson,
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept', description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nImage upscale or vary\n
-    Image upscale or vary
-    Arguments:
-        req {ImgUpscaleOrVaryRequestJson} -- Image upscale or vary request
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-            Response -- img_generate_responses    
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    req.input_image = base64_to_stream(req.input_image)
-
-    default_image_prompt = ImagePrompt(cn_img=None)
-    image_prompts_files: List[ImagePrompt] = []
-    for image_prompt in req.image_prompts:
-        image_prompt.cn_img = base64_to_stream(image_prompt.cn_img)
-        image = ImagePrompt(
-            cn_img=image_prompt.cn_img,
-            cn_stop=image_prompt.cn_stop,
-            cn_weight=image_prompt.cn_weight,
-            cn_type=image_prompt.cn_type)
-        image_prompts_files.append(image)
-    while len(image_prompts_files) <= 4:
-        image_prompts_files.append(default_image_prompt)
-    req.image_prompts = image_prompts_files
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v2/generation/image-inpaint-outpaint",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV2"])
-def img_inpaint_or_outpaint(
-    req: ImgInpaintOrOutpaintRequestJson,
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nInpaint or outpaint\n
-    Inpaint or outpaint
-    Arguments:
-        req {ImgInpaintOrOutpaintRequestJson} -- Request body
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    req.input_image = base64_to_stream(req.input_image)
-    if req.input_mask is not None:
-        req.input_mask = base64_to_stream(req.input_mask)
-    default_image_prompt = ImagePrompt(cn_img=None)
-    image_prompts_files: List[ImagePrompt] = []
-    for image_prompt in req.image_prompts:
-        image_prompt.cn_img = base64_to_stream(image_prompt.cn_img)
-        image = ImagePrompt(
-            cn_img=image_prompt.cn_img,
-            cn_stop=image_prompt.cn_stop,
-            cn_weight=image_prompt.cn_weight,
-            cn_type=image_prompt.cn_type)
-        image_prompts_files.append(image)
-    while len(image_prompts_files) <= 4:
-        image_prompts_files.append(default_image_prompt)
-    req.image_prompts = image_prompts_files
-
-    return call_worker(req, accept)
-
-
-@secure_router.post(
-        path="/v2/generation/image-prompt",
-        response_model=List[GeneratedImageResult] | AsyncJobResponse,
-        responses=img_generate_responses,
-        tags=["GenerateV2"])
-def img_prompt(
-    req: ImgPromptRequestJson,
-    accept: str = Header(None),
-    accept_query: str | None = Query(
-        None, alias='accept',
-        description="Parameter to override 'Accept' header, 'image/png' for output bytes")):
-    """\nImage prompt\n
-    Image prompt generation
-    Arguments:
-        req {ImgPromptRequest} -- Request body
-        accept {str} -- Accept header
-        accept_query {str} -- Parameter to override 'Accept' header, 'image/png' for output bytes
-    Returns:
-        Response -- img_generate_responses
-    """
-    if accept_query is not None and len(accept_query) > 0:
-        accept = accept_query
-
-    if req.input_image is not None:
-        req.input_image = base64_to_stream(req.input_image)
-    if req.input_mask is not None:
-        req.input_mask = base64_to_stream(req.input_mask)
-
-    default_image_prompt = ImagePrompt(cn_img=None)
-    image_prompts_files: List[ImagePrompt] = []
-    for image_prompt in req.image_prompts:
-        image_prompt.cn_img = base64_to_stream(image_prompt.cn_img)
-        image = ImagePrompt(
-            cn_img=image_prompt.cn_img,
-            cn_stop=image_prompt.cn_stop,
-            cn_weight=image_prompt.cn_weight,
-            cn_type=image_prompt.cn_type)
-        image_prompts_files.append(image)
-
-    while len(image_prompts_files) <= 4:
-        image_prompts_files.append(default_image_prompt)
-
-    req.image_prompts = image_prompts_files
-
-    return call_worker(req, accept)

fooocusapi/routes/query.py

@@ -1,135 +0,0 @@
-"""Query API"""
-from typing import List
-from fastapi import Depends, Response, APIRouter
-
-from fooocusapi.args import args
-
-from fooocusapi.models.common.requests import QueryJobRequest
-from fooocusapi.models.common.response import (
-    AsyncJobResponse,
-    JobHistoryInfo,
-    JobQueueInfo,
-    JobHistoryResponse,
-    AllModelNamesResponse
-)
-from fooocusapi.models.common.task import AsyncJobStage
-
-from fooocusapi.utils.api_utils import generate_async_output, api_key_auth
-from fooocusapi.task_queue import TaskType
-from fooocusapi.worker import worker_queue
-
-secure_router = APIRouter(dependencies=[Depends(api_key_auth)])
-
-
-@secure_router.get(path="/", tags=['Query'])
-def home():
-    """Home page"""
-    return Response(
-        content='Swagger-UI to: <a href="/docs">/docs</a>',
-        media_type="text/html"
-    )
-
-
-@secure_router.get(
-        path="/ping",
-        description="Returns a simple 'pong'",
-        tags=['Query'])
-async def ping():
-    """\nPing\n
-    Ping page, just to check if the fastapi is up.
-    Instant return correct, does not mean the service is available.
-    Returns:
-        A simple string pong
-    """
-    return 'pong'
-
-
-@secure_router.get(
-        path="/v1/generation/query-job",
-        response_model=AsyncJobResponse,
-        description="Query async generation job",
-        tags=['Query'])
-def query_job(req: QueryJobRequest = Depends()):
-    """query job info by id"""
-    queue_task = worker_queue.get_task(req.job_id, True)
-    if queue_task is None:
-        result = AsyncJobResponse(
-            job_id="",
-            job_type=TaskType.not_found,
-            job_stage=AsyncJobStage.error,
-            job_progress=0,
-            job_status="Job not found")
-        content = result.model_dump_json()
-        return Response(content=content, media_type='application/json', status_code=404)
-    return generate_async_output(queue_task, req.require_step_preview)
-
-
-@secure_router.get(
-        path="/v1/generation/job-queue",
-        response_model=JobQueueInfo,
-        description="Query job queue info",
-        tags=['Query'])
-def job_queue():
-    """Query job queue info"""
-    queue = JobQueueInfo(
-        running_size=len(worker_queue.queue),
-        finished_size=len(worker_queue.history),
-        last_job_id=worker_queue.last_job_id
-    )
-    return queue
-
-
-@secure_router.get(
-        path="/v1/generation/job-history",
-        response_model=JobHistoryResponse | dict,
-        description="Query historical job data",
-        tags=["Query"])
-def get_history(job_id: str = None, page: int = 0, page_size: int = 20):
-    """Fetch and return the historical tasks"""
-    queue = [
-        JobHistoryInfo(
-            job_id=item.job_id,
-            is_finished=item.is_finished
-        ) for item in worker_queue.queue
-    ]
-    if not args.persistent:
-        history = [
-            JobHistoryInfo(
-                job_id=item.job_id,
-                is_finished=item.is_finished
-            ) for item in worker_queue.history
-        ]
-        return JobHistoryResponse(history=history, queue=queue)
-
-    from fooocusapi.sql_client import query_history
-    history = query_history(task_id=job_id, page=page, page_size=page_size)
-    return {
-        "history": history,
-        "queue": queue
-    }
-
-
-@secure_router.get(
-        path="/v1/engines/all-models",
-        response_model=AllModelNamesResponse,
-        description="Get all filenames of base model and lora",
-        tags=["Query"])
-def all_models():
-    """Refresh and return all models"""
-    from modules import config
-    config.update_files()
-    models = AllModelNamesResponse(
-        model_filenames=config.model_filenames,
-        lora_filenames=config.lora_filenames)
-    return models
-
-
-@secure_router.get(
-        path="/v1/engines/styles",
-        response_model=List[str],
-        description="Get all legal Fooocus styles",
-        tags=['Query'])
-def all_styles():
-    """Return all available styles"""
-    from modules.sdxl_styles import legal_style_names
-    return legal_style_names

fooocusapi/sql_client.py

@@ -1,269 +0,0 @@
-"""
-SQLite client for Fooocus API
-"""
-import os
-import time
-import platform
-from datetime import datetime
-from typing import Optional
-import copy
-
-from sqlalchemy import Integer, Float, VARCHAR, Boolean, JSON, Text, create_engine
-from sqlalchemy.orm import declarative_base, Session, Mapped, mapped_column
-
-
-Base = declarative_base()
-
-
-if platform.system().lower() == "windows":
-    default_sqlite_db_path = os.path.join(
-        os.path.dirname(__file__), "../database.db"
-    ).replace("\\", "/")
-else:
-    default_sqlite_db_path = os.path.join(os.path.dirname(__file__), "../database.db")
-
-connection_uri = os.environ.get(
-    "FOOOCUS_DB_CONF", f"sqlite:///{default_sqlite_db_path}"
-)
-
-
-class GenerateRecord(Base):
-    """
-    GenerateRecord
-
-    __tablename__ = 'generate_record'
-    """
-
-    __tablename__ = "generate_record"
-
-    task_id: Mapped[str] = mapped_column(VARCHAR(255), nullable=False, primary_key=True)
-    task_type: Mapped[str] = mapped_column(Text, nullable=False)
-    result_url: Mapped[str] = mapped_column(Text, nullable=True)
-    finish_reason: Mapped[str] = mapped_column(Text, nullable=True)
-    date_time: Mapped[int] = mapped_column(Integer, nullable=False)
-
-    prompt: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    negative_prompt: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    style_selections: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
-    performance_selection: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    aspect_ratios_selection: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    base_model_name: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    refiner_model_name: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    refiner_switch: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-    loras: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
-    image_number: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    image_seed: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    sharpness: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-    guidance_scale: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-    advanced_params: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
-
-    input_image: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    input_mask: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    image_prompts: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
-    inpaint_additional_prompt: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    outpaint_selections: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
-    outpaint_distance_left: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    outpaint_distance_right: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    outpaint_distance_top: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    outpaint_distance_bottom: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-    uov_method: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    upscale_value: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-
-    webhook_url: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
-    require_base64: Mapped[Optional[bool]] = mapped_column(Boolean, nullable=True)
-    async_process: Mapped[Optional[bool]] = mapped_column(Boolean, nullable=True)
-
-    def __repr__(self) -> str:
-        return f"GenerateRecord(task_id={self.task_id!r}, task_type={self.task_type!r}, \
-                result_url={self.result_url!r}, finish_reason={self.finish_reason!r}, date_time={self.date_time!r}, \
-                prompt={self.prompt!r}, negative_prompt={self.negative_prompt!r}, style_selections={self.style_selections!r}, performance_selection={self.performance_selection!r}, \
-                aspect_ratios_selection={self.aspect_ratios_selection!r}, base_model_name={self.base_model_name!r}, \
-                refiner_model_name={self.refiner_model_name!r}, refiner_switch={self.refiner_switch!r}, loras={self.loras!r}, \
-                image_number={self.image_number!r}, image_seed={self.image_seed!r}, sharpness={self.sharpness!r}, \
-                guidance_scale={self.guidance_scale!r}, advanced_params={self.advanced_params!r}, input_image={self.input_image!r}, \
-                input_mask={self.input_mask!r}, image_prompts={self.image_prompts!r}, inpaint_additional_prompt={self.inpaint_additional_prompt!r}, \
-                outpaint_selections={self.outpaint_selections!r}, outpaint_distance_left={self.outpaint_distance_left!r}, outpaint_distance_right={self.outpaint_distance_right!r}, \
-                outpaint_distance_top={self.outpaint_distance_top!r}, outpaint_distance_bottom={self.outpaint_distance_bottom!r}, uov_method={self.uov_method!r}, \
-                upscale_value={self.upscale_value!r}, webhook_url={self.webhook_url!r}, require_base64={self.require_base64!r}, \
-                async_process={self.async_process!r})"
-
-
-engine = create_engine(connection_uri)
-
-session = Session(engine)
-Base.metadata.create_all(engine, checkfirst=True)
-session.close()
-
-
-def convert_to_dict_list(obj_list: list[object]) -> list[dict]:
-    """
-    Convert a list of objects to a list of dictionaries.
-    Args:
-        obj_list:
-
-    Returns:
-        dict_list:
-    """
-    dict_list = []
-    for obj in obj_list:
-        # 将对象属性转化为字典键值对
-        dict_obj = {}
-        for attr, value in vars(obj).items():
-            if (
-                not callable(value)
-                and not attr.startswith("__")
-                and not attr.startswith("_")
-            ):
-                dict_obj[attr] = value
-        task_info = {
-            "task_id": obj.task_id,
-            "task_type": obj.task_type,
-            "result_url": obj.result_url,
-            "finish_reason": obj.finish_reason,
-            "date_time": datetime.fromtimestamp(obj.date_time).strftime(
-                "%Y-%m-%d %H:%M:%S"
-            ),
-        }
-        del dict_obj["task_id"]
-        del dict_obj["task_type"]
-        del dict_obj["result_url"]
-        del dict_obj["finish_reason"]
-        del dict_obj["date_time"]
-        dict_list.append({"params": dict_obj, "task_info": task_info})
-    return dict_list
-
-
-class MySQLAlchemy:
-    """
-    MySQLAlchemy, a toolkit for managing SQLAlchemy connections and sessions.
-
-    :param uri: SQLAlchemy connection URI
-    """
-
-    def __init__(self, uri: str):
-        # 'mysql+pymysql://{username}:{password}@{host}:{port}/{database}'
-        self.engine = create_engine(uri)
-        self.session = Session(self.engine)
-
-    def store_history(self, record: dict) -> None:
-        """
-        Store history to database
-        :param record:
-        :return:
-        """
-        self.session.add_all([GenerateRecord(**record)])
-        self.session.commit()
-
-    def get_history(
-        self,
-        task_id: str = None,
-        page: int = 0,
-        page_size: int = 20,
-        order_by: str = "date_time",
-    ) -> list:
-        """
-        Get history from database
-        :param task_id:
-        :param page:
-        :param page_size:
-        :param order_by:
-        :return:
-        """
-        if task_id is not None:
-            res = (
-                self.session.query(GenerateRecord)
-                .filter(GenerateRecord.task_id == task_id)
-                .all()
-            )
-            if len(res) == 0:
-                return []
-            return convert_to_dict_list(res)
-
-        res = (
-            self.session.query(GenerateRecord)
-            .order_by(getattr(GenerateRecord, order_by).desc())
-            .offset(page * page_size)
-            .limit(page_size)
-            .all()
-        )
-        if len(res) == 0:
-            return []
-        return convert_to_dict_list(res)
-
-
-db = MySQLAlchemy(uri=connection_uri)
-
-
-def req_to_dict(req: dict) -> dict:
-    """
-    Convert request to dictionary
-    Args:
-        req:
-
-    Returns:
-
-    """
-    req["loras"] = [{"model_name": lora[0], "weight": lora[1]} for lora in req["loras"]]
-    # req["advanced_params"] = dict(zip(adv_params_keys, req["advanced_params"]))
-    req["image_prompts"] = [
-        {"cn_img": "", "cn_stop": image[1], "cn_weight": image[2], "cn_type": image[3]}
-        for image in req["image_prompts"]
-    ]
-    del req["inpaint_input_image"]
-    del req["uov_input_image"]
-    return req
-
-
-def add_history(
-    params: dict, task_type: str, task_id: str, result_url: str, finish_reason: str
-) -> None:
-    """
-    Store history to database
-    Args:
-        params:
-        task_type:
-        task_id:
-        result_url:
-        finish_reason:
-
-    Returns:
-
-    """
-    adv = copy.deepcopy(params["advanced_params"])
-    params["advanced_params"] = adv.__dict__
-    params["date_time"] = int(time.time())
-    params["task_type"] = task_type
-    params["task_id"] = task_id
-    params["result_url"] = result_url
-    params["finish_reason"] = finish_reason
-
-    del params["inpaint_input_image"]
-    del params["uov_input_image"]
-    del params["save_extension"]
-    del params["save_meta"]
-    del params["save_name"]
-    del params["meta_scheme"]
-
-    db.store_history(params)
-
-
-def query_history(
-        task_id: str = None,
-        page: int = 0,
-        page_size: int = 20,
-        order_by: str = "date_time"
-) -> list:
-    """
-    Query history from database
-    Args:
-        task_id:
-        page:
-        page_size:
-        order_by:
-
-    Returns:
-
-    """
-    return db.get_history(
-        task_id=task_id, page=page, page_size=page_size, order_by=order_by
-    )

fooocusapi/task_queue.py

@@ -1,323 +0,0 @@
-"""
-Task queue management
-
-This module provides classes and functions for managing the task queue.
-
-Classes:
-    QueueTask: A class representing a task in the queue.
-    TaskQueue: A class for managing the task queue.
-"""
-import uuid
-import time
-from typing import List, Tuple
-import numpy as np
-import requests
-
-from fooocusapi.utils.file_utils import delete_output_file, get_file_serve_url
-from fooocusapi.utils.img_utils import narray_to_base64img
-from fooocusapi.utils.logger import logger
-
-from fooocusapi.models.common.task import ImageGenerationResult, GenerationFinishReason
-from fooocusapi.parameters import ImageGenerationParams
-from fooocusapi.models.common.task import TaskType
-
-
-class QueueTask:
-    """
-    A class representing a task in the queue.
-
-    Attributes:
-        job_id (str): The unique identifier for the task, generated by uuid.
-        task_type (TaskType): The type of task.
-        is_finished (bool): Indicates whether the task has been completed.
-        finish_progress (int): The progress of the task completion.
-        in_queue_mills (int): The time the task was added to the queue, in milliseconds.
-        start_mills (int): The time the task started, in milliseconds.
-        finish_mills (int): The time the task finished, in milliseconds.
-        finish_with_error (bool): Indicates whether the task finished with an error.
-        task_status (str): The status of the task.
-        task_step_preview (str): A list of step previews for the task.
-        task_result (List[ImageGenerationResult]): The result of the task.
-        error_message (str): The error message, if any.
-        webhook_url (str): The webhook URL, if any.
-    """
-
-    job_id: str
-    task_type: TaskType
-    req_param: ImageGenerationParams
-    is_finished: bool = False
-    finish_progress: int = 0
-    in_queue_mills: int
-    start_mills: int = 0
-    finish_mills: int = 0
-    finish_with_error: bool = False
-    task_status: str | None = None
-    task_step_preview: str | None = None
-    task_result: List[ImageGenerationResult] = None
-    error_message: str | None = None
-    webhook_url: str | None = None  # attribute for individual webhook_url
-
-    def __init__(
-        self,
-        job_id: str,
-        task_type: TaskType,
-        req_param: ImageGenerationParams,
-        webhook_url: str | None = None,
-    ):
-        self.job_id = job_id
-        self.task_type = task_type
-        self.req_param = req_param
-        self.in_queue_mills = int(round(time.time() * 1000))
-        self.webhook_url = webhook_url
-
-    def set_progress(self, progress: int, status: str | None):
-        """
-        Set progress and status
-        Arguments:
-            progress {int} -- progress
-            status {str} -- status
-        """
-        progress = min(progress, 100)
-        self.finish_progress = progress
-        self.task_status = status
-
-    def set_step_preview(self, task_step_preview: str | None):
-        """set step preview
-        Set step preview
-        Arguments:
-            task_step_preview {str} -- step preview
-        """
-        self.task_step_preview = task_step_preview
-
-    def set_result(
-        self,
-        task_result: List[ImageGenerationResult],
-        finish_with_error: bool,
-        error_message: str | None = None,
-    ):
-        """set result
-        Set task result
-        Arguments:
-            task_result {List[ImageGenerationResult]} -- task result
-            finish_with_error {bool} -- finish with error
-            error_message {str} -- error message
-        """
-        if not finish_with_error:
-            self.finish_progress = 100
-            self.task_status = "Finished"
-        self.task_result = task_result
-        self.finish_with_error = finish_with_error
-        self.error_message = error_message
-
-    def __str__(self) -> str:
-        return f"QueueTask(job_id={self.job_id}, task_type={self.task_type},\
-                is_finished={self.is_finished}, finished_progress={self.finish_progress}, \
-                in_queue_mills={self.in_queue_mills}, start_mills={self.start_mills}, \
-                finish_mills={self.finish_mills}, finish_with_error={self.finish_with_error}, \
-                error_message={self.error_message}, task_status={self.task_status}, \
-                task_step_preview={self.task_step_preview}, webhook_url={self.webhook_url})"
-
-
-class TaskQueue:
-    """
-    TaskQueue is a queue of tasks that are waiting to be processed.
-
-    Attributes:
-        queue: List[QueueTask]
-        history: List[QueueTask]
-        last_job_id: str
-        webhook_url: str
-        persistent: bool
-    """
-
-    queue: List[QueueTask] = []
-    history: List[QueueTask] = []
-    last_job_id: str = None
-    webhook_url: str | None = None
-    persistent: bool = False
-
-    def __init__(
-        self,
-        queue_size: int,
-        history_size: int,
-        webhook_url: str | None = None,
-        persistent: bool | None = False,
-    ):
-        self.queue_size = queue_size
-        self.history_size = history_size
-        self.webhook_url = webhook_url
-        self.persistent = False if persistent is None else persistent
-
-    def add_task(
-        self,
-        task_type: TaskType,
-        req_param: ImageGenerationParams,
-        webhook_url: str | None = None,
-    ) -> QueueTask | None:
-        """
-        Create and add task to queue
-        :param task_type: task type
-        :param req_param: request parameters
-        :param webhook_url: webhook url
-        :returns: The created task's job_id, or None if reach the queue size limit
-        """
-        if len(self.queue) >= self.queue_size:
-            return None
-
-        if isinstance(req_param, dict):
-            req_param = ImageGenerationParams(**req_param)
-
-        job_id = str(uuid.uuid4())
-        task = QueueTask(
-            job_id=job_id,
-            task_type=task_type,
-            req_param=req_param,
-            webhook_url=webhook_url,
-        )
-        self.queue.append(task)
-        self.last_job_id = job_id
-        return task
-
-    def get_task(self, job_id: str, include_history: bool = False) -> QueueTask | None:
-        """
-        Get task by job_id
-        :param job_id: job id
-        :param include_history: whether to include history tasks
-        :returns: The task with the given job_id, or None if not found
-        """
-        for task in self.queue:
-            if task.job_id == job_id:
-                return task
-
-        if include_history:
-            for task in self.history:
-                if task.job_id == job_id:
-                    return task
-
-        return None
-
-    def is_task_ready_to_start(self, job_id: str) -> bool:
-        """
-        Check if the task is ready to start
-        :param job_id: job id
-        :returns: True if the task is ready to start, False otherwise
-        """
-        task = self.get_task(job_id)
-        if task is None:
-            return False
-
-        return self.queue[0].job_id == job_id
-
-    def is_task_finished(self, job_id: str) -> bool:
-        """
-        Check if the task is finished
-        :param job_id: job id
-        :returns: True if the task is finished, False otherwise
-        """
-        task = self.get_task(job_id, True)
-        if task is None:
-            return False
-
-        return task.is_finished
-
-    def start_task(self, job_id: str):
-        """
-        Start task by job_id
-        :param job_id: job id
-        """
-        task = self.get_task(job_id)
-        if task is not None:
-            task.start_mills = int(round(time.time() * 1000))
-
-    def finish_task(self, job_id: str):
-        """
-        Finish task by job_id
-        :param job_id: job id
-        """
-        task = self.get_task(job_id)
-        if task is not None:
-            task.is_finished = True
-            task.finish_mills = int(round(time.time() * 1000))
-
-            # Use the task's webhook_url if available, else use the default
-            webhook_url = task.webhook_url or self.webhook_url
-
-            data = {"job_id": task.job_id, "job_result": []}
-
-            if isinstance(task.task_result, List):
-                for item in task.task_result:
-                    data["job_result"].append(
-                        {
-                            "url": get_file_serve_url(item.im) if item.im else None,
-                            "seed": item.seed if item.seed else "-1",
-                        }
-                    )
-
-            # Send webhook
-            if task.is_finished and webhook_url:
-                try:
-                    res = requests.post(webhook_url, json=data, timeout=15)
-                    print(f"Call webhook response status: {res.status_code}")
-                except Exception as e:
-                    print("Call webhook error:", e)
-
-            # Move task to history
-            self.queue.remove(task)
-            self.history.append(task)
-
-            # save history to database
-            if self.persistent:
-                from fooocusapi.sql_client import add_history
-
-                add_history(
-                    params=task.req_param.to_dict(),
-                    task_type=task.task_type.value,
-                    task_id=task.job_id,
-                    result_url=",".join([job["url"] for job in data["job_result"]]),
-                    finish_reason=task.task_result[0].finish_reason.value,
-                )
-
-            # Clean history
-            if len(self.history) > self.history_size != 0:
-                removed_task = self.history.pop(0)
-                if isinstance(removed_task.task_result, List):
-                    for item in removed_task.task_result:
-                        if (
-                            isinstance(item, ImageGenerationResult)
-                            and item.finish_reason == GenerationFinishReason.success
-                            and item.im is not None
-                        ):
-                            delete_output_file(item.im)
-                logger.std_info(
-                    f"[TaskQueue] Clean task history, remove task: {removed_task.job_id}"
-                )
-
-
-class TaskOutputs:
-    """
-    TaskOutputs is a container for task outputs
-    """
-
-    outputs = []
-
-    def __init__(self, task: QueueTask):
-        self.task = task
-
-    def append(self, args: List[any]):
-        """
-        Append output to task outputs list
-        :param args: output arguments
-        """
-        self.outputs.append(args)
-        if len(args) >= 2:
-            if (
-                args[0] == "preview"
-                and isinstance(args[1], Tuple)
-                and len(args[1]) >= 2
-            ):
-                number = args[1][0]
-                text = args[1][1]
-                self.task.set_progress(number, text)
-                if len(args[1]) >= 3 and isinstance(args[1][2], np.ndarray):
-                    base64_preview_img = narray_to_base64img(args[1][2])
-                    self.task.set_step_preview(base64_preview_img)

fooocusapi/utils/api_utils.py

@@ -1,291 +0,0 @@
-"""some utils for api"""
-from typing import List
-
-from fastapi import Response
-from fastapi.security import APIKeyHeader
-from fastapi import HTTPException, Security
-
-from modules import flags
-from modules import config
-from modules.sdxl_styles import legal_style_names
-
-from fooocusapi.args import args
-from fooocusapi.utils.img_utils import read_input_image
-from fooocusapi.utils.file_utils import (
-    get_file_serve_url,
-    output_file_to_base64img,
-    output_file_to_bytesimg
-)
-from fooocusapi.utils.logger import logger
-from fooocusapi.models.common.requests import (
-    CommonRequest as Text2ImgRequest
-)
-from fooocusapi.models.common.response import (
-    AsyncJobResponse,
-    AsyncJobStage,
-    GeneratedImageResult
-)
-from fooocusapi.models.requests_v1 import (
-    ImgInpaintOrOutpaintRequest,
-    ImgPromptRequest,
-    ImgUpscaleOrVaryRequest
-)
-from fooocusapi.models.requests_v2 import (
-    Text2ImgRequestWithPrompt,
-    ImgInpaintOrOutpaintRequestJson,
-    ImgUpscaleOrVaryRequestJson,
-    ImgPromptRequestJson
-)
-from fooocusapi.models.common.task import (
-    ImageGenerationResult,
-    GenerationFinishReason
-)
-from fooocusapi.configs.default import (
-    default_inpaint_engine_version,
-    default_sampler,
-    default_scheduler,
-    default_base_model_name,
-    default_refiner_model_name
-)
-
-from fooocusapi.parameters import ImageGenerationParams
-from fooocusapi.task_queue import QueueTask
-
-
-api_key_header = APIKeyHeader(name="X-API-KEY", auto_error=False)
-
-
-def api_key_auth(apikey: str = Security(api_key_header)):
-    """
-    Check if the API key is valid, API key is not required if no API key is set
-    Args:
-        apikey: API key
-    returns:
-        None if API key is not set, otherwise raise HTTPException
-    """
-    if args.apikey is None:
-        return  # Skip API key check if no API key is set
-    if apikey != args.apikey:
-        raise HTTPException(status_code=403, detail="Forbidden")
-
-
-def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams:
-    """
-    Convert Request to ImageGenerationParams
-    Args:
-        req: Request, Text2ImgRequest and classes inherited from Text2ImgRequest
-    returns:
-        ImageGenerationParams
-    """
-    config.update_files()
-    if req.base_model_name is not None:
-        if req.base_model_name not in config.model_filenames:
-            logger.std_warn(f"[Warning] Wrong base_model_name input: {req.base_model_name}, using default")
-            req.base_model_name = default_base_model_name
-
-    if req.refiner_model_name is not None and req.refiner_model_name != 'None':
-        if req.refiner_model_name not in config.model_filenames:
-            logger.std_warn(f"[Warning] Wrong refiner_model_name input: {req.refiner_model_name}, using default")
-            req.refiner_model_name = default_refiner_model_name
-
-    for lora in req.loras:
-        if lora.model_name != 'None' and lora.model_name not in config.lora_filenames:
-            logger.std_warn(f"[Warning] Wrong lora model_name input: {lora.model_name}, using 'None'")
-            lora.model_name = 'None'
-
-    prompt = req.prompt
-    negative_prompt = req.negative_prompt
-    style_selections = [
-        s for s in req.style_selections if s in legal_style_names]
-    performance_selection = req.performance_selection.value
-    aspect_ratios_selection = req.aspect_ratios_selection
-    image_number = req.image_number
-    image_seed = None if req.image_seed == -1 else req.image_seed
-    sharpness = req.sharpness
-    guidance_scale = req.guidance_scale
-    base_model_name = req.base_model_name
-    refiner_model_name = req.refiner_model_name
-    refiner_switch = req.refiner_switch
-    loras = [(lora.model_name, lora.weight) for lora in req.loras]
-    uov_input_image = None
-    if not isinstance(req, Text2ImgRequestWithPrompt):
-        if isinstance(req, (ImgUpscaleOrVaryRequest, ImgUpscaleOrVaryRequestJson)):
-            uov_input_image = read_input_image(req.input_image)
-    uov_method = flags.disabled if not isinstance(req, (ImgUpscaleOrVaryRequest, ImgUpscaleOrVaryRequestJson)) else req.uov_method.value
-    upscale_value = None if not isinstance(req, (ImgUpscaleOrVaryRequest, ImgUpscaleOrVaryRequestJson)) else req.upscale_value
-    outpaint_selections = [] if not isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) else [
-        s.value for s in req.outpaint_selections]
-    outpaint_distance_left = None if not isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) else req.outpaint_distance_left
-    outpaint_distance_right = None if not isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) else req.outpaint_distance_right
-    outpaint_distance_top = None if not isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) else req.outpaint_distance_top
-    outpaint_distance_bottom = None if not isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) else req.outpaint_distance_bottom
-
-    if refiner_model_name == '':
-        refiner_model_name = 'None'
-
-    inpaint_input_image = None
-    inpaint_additional_prompt = None
-    if isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)) and req.input_image is not None:
-        inpaint_additional_prompt = req.inpaint_additional_prompt
-        input_image = read_input_image(req.input_image)
-        input_mask = None
-        if req.input_mask is not None:
-            input_mask = read_input_image(req.input_mask)
-        inpaint_input_image = {
-            'image': input_image,
-            'mask': input_mask
-        }
-
-    image_prompts = []
-    if isinstance(req, (ImgInpaintOrOutpaintRequestJson, ImgPromptRequest, ImgPromptRequestJson, ImgUpscaleOrVaryRequestJson, Text2ImgRequestWithPrompt)):
-        # Auto set mixing_image_prompt_and_inpaint to True
-        if len(req.image_prompts) > 0 and uov_input_image is not None:
-            print("[INFO] Mixing image prompt and vary upscale is set to True")
-            req.advanced_params.mixing_image_prompt_and_vary_upscale = True
-        elif len(req.image_prompts) > 0 and not isinstance(req, Text2ImgRequestWithPrompt) and req.input_image is not None:
-            print("[INFO] Mixing image prompt and inpaint is set to True")
-            req.advanced_params.mixing_image_prompt_and_inpaint = True
-
-        for img_prompt in req.image_prompts:
-            if img_prompt.cn_img is not None:
-                cn_img = read_input_image(img_prompt.cn_img)
-                if img_prompt.cn_stop is None or img_prompt.cn_stop == 0:
-                    img_prompt.cn_stop = flags.default_parameters[img_prompt.cn_type.value][0]
-                if img_prompt.cn_weight is None or img_prompt.cn_weight == 0:
-                    img_prompt.cn_weight = flags.default_parameters[img_prompt.cn_type.value][1]
-                image_prompts.append(
-                    (cn_img, img_prompt.cn_stop, img_prompt.cn_weight, img_prompt.cn_type.value))
-
-    advanced_params = None
-    if req.advanced_params is not None:
-        adp = req.advanced_params
-
-        if adp.refiner_swap_method not in ['joint', 'separate', 'vae']:
-            print(f"[Warning] Wrong refiner_swap_method input: {adp.refiner_swap_method}, using default")
-            adp.refiner_swap_method = 'joint'
-
-        if adp.sampler_name not in flags.sampler_list:
-            print(f"[Warning] Wrong sampler_name input: {adp.sampler_name}, using default")
-            adp.sampler_name = default_sampler
-
-        if adp.scheduler_name not in flags.scheduler_list:
-            print(f"[Warning] Wrong scheduler_name input: {adp.scheduler_name}, using default")
-            adp.scheduler_name = default_scheduler
-
-        if adp.inpaint_engine not in flags.inpaint_engine_versions:
-            print(f"[Warning] Wrong inpaint_engine input: {adp.inpaint_engine}, using default")
-            adp.inpaint_engine = default_inpaint_engine_version
-
-        advanced_params = adp
-
-    return ImageGenerationParams(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        style_selections=style_selections,
-        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=upscale_value,
-        outpaint_selections=outpaint_selections,
-        outpaint_distance_left=outpaint_distance_left,
-        outpaint_distance_right=outpaint_distance_right,
-        outpaint_distance_top=outpaint_distance_top,
-        outpaint_distance_bottom=outpaint_distance_bottom,
-        inpaint_input_image=inpaint_input_image,
-        inpaint_additional_prompt=inpaint_additional_prompt,
-        image_prompts=image_prompts,
-        advanced_params=advanced_params,
-        save_meta=req.save_meta,
-        meta_scheme=req.meta_scheme,
-        save_name=req.save_name,
-        save_extension=req.save_extension,
-        require_base64=req.require_base64,
-    )
-
-
-def generate_async_output(
-        task: QueueTask,
-        require_step_preview: bool = False) -> AsyncJobResponse:
-    """
-    Generate output for async job
-    Arguments:
-        task: QueueTask
-        require_step_preview: bool
-    Returns:
-        AsyncJobResponse
-    """
-    job_stage = AsyncJobStage.running
-    job_result = None
-
-    if task.start_mills == 0:
-        job_stage = AsyncJobStage.waiting
-
-    if task.is_finished:
-        if task.finish_with_error:
-            job_stage = AsyncJobStage.error
-        elif task.task_result is not None:
-            job_stage = AsyncJobStage.success
-            job_result = generate_image_result_output(task.task_result, task.req_param.require_base64)
-
-    result = AsyncJobResponse(
-        job_id=task.job_id,
-        job_type=task.task_type,
-        job_stage=job_stage,
-        job_progress=task.finish_progress,
-        job_status=task.task_status,
-        job_step_preview=task.task_step_preview if require_step_preview else None,
-        job_result=job_result)
-    return result
-
-
-def generate_streaming_output(results: List[ImageGenerationResult]) -> Response:
-    """
-    Generate streaming output for image generation results.
-    Args:
-        results (List[ImageGenerationResult]): List of image generation results.
-    Returns:
-        Response: Streaming response object, bytes image.
-    """
-    if len(results) == 0:
-        return Response(status_code=500)
-    result = results[0]
-    if result.finish_reason == GenerationFinishReason.queue_is_full:
-        return Response(status_code=409, content=result.finish_reason.value)
-    if result.finish_reason == GenerationFinishReason.user_cancel:
-        return Response(status_code=400, content=result.finish_reason.value)
-    if result.finish_reason == GenerationFinishReason.error:
-        return Response(status_code=500, content=result.finish_reason.value)
-
-    img_bytes = output_file_to_bytesimg(results[0].im)
-    return Response(img_bytes, media_type='image/png')
-
-
-def generate_image_result_output(
-        results: List[ImageGenerationResult],
-        require_base64: bool) -> List[GeneratedImageResult]:
-    """
-    Generate image result output
-    Arguments:
-        results: List[ImageGenerationResult]
-        require_base64: bool
-    Returns:
-        List[GeneratedImageResult]
-    """
-    results = [
-        GeneratedImageResult(
-            base64=output_file_to_base64img(item.im) if require_base64 else None,
-            url=get_file_serve_url(item.im),
-            seed=str(item.seed),
-            finish_reason=item.finish_reason
-            ) for item in results
-        ]
-    return results

fooocusapi/utils/call_worker.py

@@ -1,97 +0,0 @@
-"""function for call generate worker"""
-from typing import List
-from fastapi import Response
-
-from fooocusapi.models.common.requests import (
-    CommonRequest as Text2ImgRequest
-)
-from fooocusapi.models.common.response import (
-    AsyncJobResponse,
-    GeneratedImageResult
-)
-from fooocusapi.models.common.task import (
-    GenerationFinishReason,
-    ImageGenerationResult,
-    AsyncJobStage,
-    TaskType
-)
-from fooocusapi.utils.api_utils import (
-    req_to_params,
-    generate_async_output,
-    generate_streaming_output,
-    generate_image_result_output
-)
-from fooocusapi.models.requests_v1 import (
-    ImgUpscaleOrVaryRequest,
-    ImgPromptRequest,
-    ImgInpaintOrOutpaintRequest
-)
-from fooocusapi.models.requests_v2 import (
-    ImgInpaintOrOutpaintRequestJson,
-    ImgPromptRequestJson,
-    ImgUpscaleOrVaryRequestJson
-)
-from fooocusapi.worker import worker_queue, blocking_get_task_result
-
-
-def get_task_type(req: Text2ImgRequest) -> TaskType:
-    """return task type"""
-    if isinstance(req, (ImgUpscaleOrVaryRequest, ImgUpscaleOrVaryRequestJson)):
-        return TaskType.img_uov
-    if isinstance(req, (ImgPromptRequest, ImgPromptRequestJson)):
-        return TaskType.img_prompt
-    if isinstance(req, (ImgInpaintOrOutpaintRequest, ImgInpaintOrOutpaintRequestJson)):
-        return TaskType.img_inpaint_outpaint
-    return TaskType.text_2_img
-
-
-def call_worker(req: Text2ImgRequest, accept: str) -> Response | AsyncJobResponse | List[GeneratedImageResult]:
-    """call generation worker"""
-    if accept == 'image/png':
-        streaming_output = True
-        # image_number auto set to 1 in streaming mode
-        req.image_number = 1
-    else:
-        streaming_output = False
-
-    task_type = get_task_type(req)
-    params = req_to_params(req)
-    async_task = worker_queue.add_task(task_type, params, req.webhook_url)
-
-    if async_task is None:
-        # add to worker queue failed
-        failure_results = [
-            ImageGenerationResult(
-                im=None,
-                seed='',
-                finish_reason=GenerationFinishReason.queue_is_full
-            )]
-
-        if streaming_output:
-            return generate_streaming_output(failure_results)
-        if req.async_process:
-            return AsyncJobResponse(
-                job_id='',
-                job_type=get_task_type(req),
-                job_stage=AsyncJobStage.error,
-                job_progress=0,
-                job_status=None,
-                job_step_preview=None,
-                job_result=[GeneratedImageResult(
-                    base64=None,
-                    url=None,
-                    seed='',
-                    finish_reason=GenerationFinishReason.queue_is_full
-                )])
-        return generate_image_result_output(failure_results, False)
-
-    if req.async_process:
-        # return async response directly
-        return generate_async_output(async_task)
-
-    # blocking get generation result
-    results = blocking_get_task_result(async_task.job_id)
-
-    if streaming_output:
-        return generate_streaming_output(results)
-    return generate_image_result_output(results, req.require_base64)

fooocusapi/utils/file_utils.py

@@ -1,143 +0,0 @@
-# -*- coding: utf-8 -*-
-
-""" File utils
-
-Use for managing generated files
-
-@file: file_utils.py
-@author: Konie
-@update: 2024-03-22
-"""
-import base64
-import datetime
-from io import BytesIO
-import os
-import json
-from pathlib import Path
-import numpy as np
-from PIL import Image
-from PIL.PngImagePlugin import PngInfo
-
-from fooocusapi.utils.logger import logger
-
-
-output_dir = os.path.abspath(os.path.join(
-    os.path.dirname(__file__), '../..', 'outputs', 'files'))
-os.makedirs(output_dir, exist_ok=True)
-
-STATIC_SERVER_BASE = 'http://127.0.0.1:8888/files/'
-
-
-def save_output_file(
-        img: np.ndarray,
-        image_meta: dict = None,
-        image_name: str = '',
-        extension: str = 'png') -> str:
-    """
-    Save np image to file
-    Args:
-        img: np.ndarray image to save
-        image_meta: dict of image metadata
-        image_name: str of image name
-        extension: str of image extension
-    Returns:
-        str of file name
-    """
-    current_time = datetime.datetime.now()
-    date_string = current_time.strftime("%Y-%m-%d")
-
-    filename = os.path.join(date_string, image_name + '.' + extension)
-    file_path = os.path.join(output_dir, filename)
-
-    if extension not in ['png', 'jpg', 'webp']:
-        extension = 'png'
-    image_format = Image.registered_extensions()['.'+extension]
-
-    if image_meta is None:
-        image_meta = {}
-
-    meta = None
-    if extension == 'png'and image_meta != {}:
-        meta = PngInfo()
-        meta.add_text("parameters", json.dumps(image_meta))
-        meta.add_text("fooocus_scheme", image_meta['metadata_scheme'])
-
-    os.makedirs(os.path.dirname(file_path), exist_ok=True)
-    Image.fromarray(img).save(
-        file_path,
-        format=image_format,
-        pnginfo=meta,
-        optimize=True)
-    return Path(filename).as_posix()
-
-
-def delete_output_file(filename: str):
-    """
-    Delete files specified in the output directory
-    Args:
-        filename: str of file name
-    """
-    file_path = os.path.join(output_dir, filename)
-    if not os.path.exists(file_path) or not os.path.isfile(file_path):
-        logger.std_warn(f'[Fooocus API] {filename} not exists or is not a file')
-    try:
-        os.remove(file_path)
-        logger.std_info(f'[Fooocus API] Delete output file: {filename}')
-    except OSError:
-        logger.std_error(f'[Fooocus API] Delete output file failed: {filename}')
-
-
-def output_file_to_base64img(filename: str | None) -> str | None:
-    """
-    Convert an image file to a base64 string.
-    Args:
-        filename: str of file name
-    return: str of base64 string
-    """
-    if filename is None:
-        return None
-    file_path = os.path.join(output_dir, filename)
-    if not os.path.exists(file_path) or not os.path.isfile(file_path):
-        return None
-
-    ext = filename.split('.')[-1]
-    if ext.lower() not in ['png', 'jpg', 'webp', 'jpeg']:
-        ext = 'png'
-    img = Image.open(file_path)
-    output_buffer = BytesIO()
-    img.save(output_buffer, format=ext.upper())
-    byte_data = output_buffer.getvalue()
-    base64_str = base64.b64encode(byte_data).decode('utf-8')
-    return f"data:image/{ext};base64," + base64_str
-
-
-def output_file_to_bytesimg(filename: str | None) -> bytes | None:
-    """
-    Convert an image file to a bytes string.
-    Args:
-        filename: str of file name
-    return: bytes of image data
-    """
-    if filename is None:
-        return None
-    file_path = os.path.join(output_dir, filename)
-    if not os.path.exists(file_path) or not os.path.isfile(file_path):
-        return None
-
-    img = Image.open(file_path)
-    output_buffer = BytesIO()
-    img.save(output_buffer, format='PNG')
-    byte_data = output_buffer.getvalue()
-    return byte_data
-
-
-def get_file_serve_url(filename: str | None) -> str | None:
-    """
-    Get the static serve url of an image file.
-    Args:
-        filename: str of file name
-    return: str of static serve url
-    """
-    if filename is None:
-        return None
-    return STATIC_SERVER_BASE + filename.replace('\\', '/')

fooocusapi/utils/img_utils.py

@@ -1,198 +0,0 @@
-"""
-Image process utils. Used to verify, convert and store Images.
-
-@file: img_utils.py
-@author: Konie
-@update: 2024-03-23
-"""
-import base64
-from io import BytesIO
-from fastapi import UploadFile
-from PIL import Image
-
-import requests
-import numpy as np
-
-
-def upload2base64(image: UploadFile) -> str | None:
-    """
-    Convert UploadFile obj to base64 string
-    Args:
-        image (UploadFile): UploadFile obj
-    Returns:
-        str: base64 string, None for None
-    """
-    if image is None:
-        return None
-    image_bytes = image.file.read()
-    image_base64 = base64.b64encode(image_bytes).decode("utf-8")
-    return image_base64
-
-
-def narray_to_base64img(narray: np.ndarray) -> str | None:
-    """
-    Convert numpy array to base64 image string.
-    Args:
-        narray: numpy array
-    Returns:
-        base64 image string
-    """
-    if narray is None:
-        return None
-
-    img = Image.fromarray(narray)
-    output_buffer = BytesIO()
-    img.save(output_buffer, format='PNG')
-    byte_data = output_buffer.getvalue()
-    base64_str = base64.b64encode(byte_data).decode('utf-8')
-    return base64_str
-
-
-def narray_to_bytesimg(narray) -> bytes | None:
-    """
-    Convert numpy array to bytes image.
-    Args:
-        narray: numpy array
-    Returns:
-        bytes image
-    """
-    if narray is None:
-        return None
-
-    img = Image.fromarray(narray)
-    output_buffer = BytesIO()
-    img.save(output_buffer, format='PNG')
-    byte_data = output_buffer.getvalue()
-    return byte_data
-
-
-def read_input_image(input_image: UploadFile | str | None) -> np.ndarray | None:
-    """
-    Read input image from UploadFile or base64 string.
-    Args:
-        input_image: UploadFile, or base64 image string, or None
-    Returns:
-        numpy array of image
-    """
-    if input_image is None or input_image == '':
-        return None
-    if isinstance(input_image, str):
-        input_image_bytes = base64.b64decode(input_image)
-    else:
-        input_image_bytes = input_image.file.read()
-    pil_image = Image.open(BytesIO(input_image_bytes))
-    image = np.array(pil_image)
-    return image
-
-
-def base64_to_stream(image: str) -> UploadFile | None:
-    """
-    Convert base64 image string to UploadFile.
-    Args:
-        image: base64 image string
-    Returns:
-        UploadFile or None
-    """
-    if image in ['', None, 'None', 'none', 'string', 'null']:
-        return None
-    if image.startswith('http'):
-        return get_check_image(url=image)
-    if image.startswith('data:image'):
-        image = image.split(sep=',', maxsplit=1)[1]
-    image_bytes = base64.b64decode(image)
-    byte_stream = BytesIO()
-    byte_stream.write(image_bytes)
-    byte_stream.seek(0)
-    return UploadFile(file=byte_stream)
-
-
-def get_check_image(url: str) -> UploadFile | None:
-    """
-    Get image from url and check if it's valid.
-    Args:
-        url: image url
-    Returns:
-        UploadFile or None
-    """
-    if url == '':
-        return None
-    headers = {
-        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
-    }
-    try:
-        response = requests.get(url, headers=headers, timeout=10)
-        binary_image = response.content
-    except Exception:
-        return None
-    try:
-        buffer = BytesIO(binary_image)
-        Image.open(buffer)  # This validates the image
-    except Exception:
-        return None
-    byte_stream = BytesIO()
-    byte_stream.write(binary_image)
-    byte_stream.seek(0)
-    return UploadFile(file=byte_stream)
-
-
-def bytes_image_to_io(binary_image: bytes) -> BytesIO | None:
-    """
-    Convert bytes image to BytesIO.
-    Args:
-        binary_image: bytes image
-    Returns:
-        BytesIO or None
-    """
-    try:
-        buffer = BytesIO(binary_image)
-        Image.open(buffer)
-    except Exception:
-        return None
-    byte_stream = BytesIO()
-    byte_stream.write(binary_image)
-    byte_stream.seek(0)
-    return byte_stream
-
-
-def bytes_to_base64img(byte_data: bytes) -> str | None:
-    """
-    Convert bytes image to base64 image string.
-    Args:
-        byte_data: bytes image
-    Returns:
-        base64 image string or None
-    """
-    if byte_data is None:
-        return None
-
-    base64_str = base64.b64encode(byte_data).decode('utf-8')
-    return base64_str
-
-
-def base64_to_bytesimg(base64_str: str) -> bytes | None:
-    """
-    Convert base64 image string to bytes image.
-    Args:
-        base64_str: base64 image string
-    Returns:
-        bytes image or None
-    """
-    if base64_str == '':
-        return None
-    bytes_image = base64.b64decode(base64_str)
-    return bytes_image
-
-
-def base64_to_narray(base64_str: str) -> np.ndarray | None:
-    """
-    Convert base64 image string to numpy array.
-    Args:
-        base64_str: base64 image string
-    Returns:
-        numpy array or None
-    """
-    if base64_str == '':
-        return None
-    bytes_image = base64.b64decode(base64_str)
-    image = np.frombuffer(bytes_image, np.uint8)
-    return image

fooocusapi/utils/logger.py

@@ -1,132 +0,0 @@
-# -*- coding: utf-8 -*-
-
-""" A simply logger.
-
-This module is used to log the program.
-
-@file: logger.py
-@author: mrhan1993
-@update: 2024-03-22
-"""
-import logging
-import os
-import sys
-
-try:
-    from colorlog import ColoredFormatter
-except ImportError:
-    from fooocusapi.utils.tools import run_pip
-    run_pip(
-        command="install colorlog",
-        desc="Install colorlog for logger.",
-        live=True
-    )
-finally:
-    from colorlog import ColoredFormatter
-
-
-own_path = os.path.dirname(os.path.abspath(__file__))
-log_dir = "logs"
-default_log_path = os.path.join(own_path, '../../', log_dir)
-
-std_formatter = ColoredFormatter(
-    fmt="%(log_color)s[%(asctime)s] %(levelname)-8s%(reset)s %(blue)s%(message)s",
-    datefmt='%Y-%m-%d %H:%M:%S',
-    reset=True,
-    log_colors={
-        'DEBUG': 'cyan',
-        'INFO': 'green',
-        'WARNING': 'yellow',
-        'ERROR': 'red',
-        'CRITICAL': 'red,bg_white',
-    },
-    secondary_log_colors={},
-    style='%'
-)
-
-file_formatter = ColoredFormatter(
-    fmt="[%(asctime)s] %(levelname)-8s%(reset)s %(message)s",
-    datefmt='%Y-%m-%d %H:%M:%S',
-    reset=True,
-    no_color=True,
-    style='%'
-)
-
-
-class ConfigLogger:
-    """
-    Configure logger.
-    :param log_path: log file path, better absolute path
-    :param std_format: stdout log format
-    :param file_format: file log format
-    """
-    def __init__(self,
-                 log_path: str = default_log_path,
-                 std_format: ColoredFormatter = std_formatter,
-                 file_format: ColoredFormatter = file_formatter) -> None:
-        self.log_path = log_path
-        self.std_format = std_format
-        self.file_format = file_format
-
-
-class Logger:
-    """
-    A simple logger.
-    :param log_name: log name
-    :param config: config logger
-    """
-    def __init__(self, log_name, config: ConfigLogger = ConfigLogger()):
-        log_path = config.log_path
-        err_log_path = os.path.join(str(log_path), f"{log_name}_error.log")
-        info_log_path = os.path.join(str(log_path), f"{log_name}_info.log")
-        if not os.path.exists(log_path):
-            os.makedirs(log_path, exist_ok=True)
-
-        self._file_logger = logging.getLogger(log_name)
-        self._file_logger.setLevel("INFO")
-
-        self._std_logger = logging.getLogger()
-        self._std_logger.setLevel("INFO")
-
-        # 创建一个ERROR级别的handler，将日志记录到error.log文件中
-        error_handler = logging.FileHandler(err_log_path, encoding='utf-8')
-        error_handler.setLevel(logging.ERROR)
-
-        # 创建一个INFO级别的handler，将日志记录到info.log文件中
-        info_handler = logging.FileHandler(info_log_path, encoding='utf-8')
-        info_handler.setLevel(logging.INFO)
-
-        # 创建一个 stream handler
-        stream_handler = logging.StreamHandler(sys.stdout)
-
-        error_handler.setFormatter(config.file_format)
-        info_handler.setFormatter(config.file_format)
-        stream_handler.setFormatter(config.std_format)
-
-        # 将handler添加到logger中
-        self._file_logger.addHandler(error_handler)
-        self._file_logger.addHandler(info_handler)
-        self._std_logger.addHandler(stream_handler)
-
-    def file_error(self, message):
-        """file error log"""
-        self._file_logger.error(message)
-
-    def file_info(self, message):
-        """file info log"""
-        self._file_logger.info(message)
-
-    def std_info(self, message):
-        """std info log"""
-        self._std_logger.info(message)
-
-    def std_warn(self, message):
-        """std warn log"""
-        self._std_logger.warning(message)
-
-    def std_error(self, message):
-        """std error log"""
-        self._std_logger.error(message)
-
-
-logger = Logger(log_name="fooocus_api")

fooocusapi/utils/lora_manager.py

@@ -1,71 +0,0 @@
-import hashlib
-import os
-import requests
-import tarfile
-
-def _hash_url(url):
-    """Generates a hash value for a given URL."""
-    return hashlib.md5(url.encode('utf-8')).hexdigest()
-
-class LoraManager:
-    """
-    Manager loras from url
-    """
-    def __init__(self):
-        self.cache_dir = os.path.join(
-            os.path.dirname(os.path.realpath(__file__)),
-            '../../',
-            'repositories/Fooocus/models/loras')
-
-    def _download_lora(self, url):
-        """
-        Downloads a LoRa from a URL, saves it in the cache, and if it's a .tar file, extracts it and returns the .safetensors file.
-        """
-        url_hash = _hash_url(url)
-        file_ext = url.split('.')[-1]
-        filepath = os.path.join(self.cache_dir, f"{url_hash}.{file_ext}")
-
-        if not os.path.exists(filepath):
-            print(f"Start download for: {url}")
-
-            try:
-                response = requests.get(url, timeout=10, stream=True)
-                response.raise_for_status()
-                with open(filepath, 'wb') as f:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        f.write(chunk)
-                
-                if file_ext == "tar":
-                    print("Extracting the tar file...")
-                    with tarfile.open(filepath, 'r:*') as tar:
-                        tar.extractall(path=self.cache_dir)
-                    print("Extraction completed.")
-                    return self._find_safetensors_file(self.cache_dir)
-
-                print(f"Download successfully, saved as {filepath}")
-            except Exception as e:
-                raise Exception(f"Error downloading {url}: {e}") from e
-
-        else:
-            print(f"LoRa already downloaded {url}")
-
-        return filepath
-
-    def _find_safetensors_file(self, directory):
-        """
-        Finds the first .safetensors file in the specified directory.
-        """
-        print("Searching for .safetensors file.")
-        for root, dirs, files in os.walk(directory):
-            for file in files:
-                if file.endswith('.safetensors'):
-                    return os.path.join(root, file)
-        raise FileNotFoundError("No .safetensors file found in the extracted files.")
-
-    def check(self, urls):
-        """Manages the specified LoRAs: downloads missing ones and returns their file names."""
-        paths = []
-        for url in urls:
-            path = self._download_lora(url)
-            paths.append(path)
-        return paths

fooocusapi/utils/model_loader.py

@@ -1,46 +0,0 @@
-# -*- coding: utf-8 -*-
-
-"""
-Download models from url
-
-@file: model_loader.py
-@author: Konie
-@update: 2024-03-22 
-"""
-from modules.model_loader import load_file_from_url
-
-
-def download_models():
-    """
-    Download models from config
-    """
-    vae_approx_filenames = [
-        ('xlvaeapp.pth', 'https://huggingface.co/lllyasviel/misc/resolve/main/xlvaeapp.pth'),
-        ('vaeapp_sd15.pth', 'https://huggingface.co/lllyasviel/misc/resolve/main/vaeapp_sd15.pt'),
-        ('xl-to-v1_interposer-v3.1.safetensors', 'https://huggingface.co/lllyasviel/misc/resolve/main/xl-to-v1_interposer-v3.1.safetensors')
-    ]
-
-    from modules.config import (
-        paths_checkpoints as modelfile_path,
-        paths_loras as lorafile_path,
-        path_vae_approx as vae_approx_path,
-        path_fooocus_expansion as fooocus_expansion_path,
-        path_embeddings as embeddings_path,
-        checkpoint_downloads,
-        embeddings_downloads,
-        lora_downloads)
-
-    for file_name, url in checkpoint_downloads.items():
-        load_file_from_url(url=url, model_dir=modelfile_path[0], file_name=file_name)
-    for file_name, url in embeddings_downloads.items():
-        load_file_from_url(url=url, model_dir=embeddings_path, file_name=file_name)
-    for file_name, url in lora_downloads.items():
-        load_file_from_url(url=url, model_dir=lorafile_path[0], file_name=file_name)
-    for file_name, url in vae_approx_filenames:
-        load_file_from_url(url=url, model_dir=vae_approx_path, file_name=file_name)
-
-    load_file_from_url(
-        url='https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_expansion.bin',
-        model_dir=fooocus_expansion_path,
-        file_name='pytorch_model.bin'
-    )

fooocusapi/utils/tools.py

@@ -1,159 +0,0 @@
-# -*- coding: utf-8 -*-
-
-""" Some tools
-
-@file: tools.py
-@author: Konie
-@update: 2024-03-22
-"""
-# pylint: disable=line-too-long
-# pylint: disable=broad-exception-caught
-import os
-import sys
-import re
-import subprocess
-from importlib.util import find_spec
-from importlib import metadata
-from packaging import version
-
-
-PYTHON_EXEC = sys.executable
-INDEX_URL = os.environ.get('INDEX_URL', "")
-PATTERN = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
-
-
-# This function was copied from [Fooocus](https://github.com/lllyasviel/Fooocus) repository.
-def run_command(command: str,
-                desc: str = None,
-                error_desc: str = None,
-                custom_env: str = None,
-                live: bool = True) -> str:
-    """
-    Run a command and return the output
-    Args:
-        command: Command to run
-        desc: Description of the command
-        error_desc: Description of the error
-        custom_env: Custom environment variables
-        live: Whether to print the output
-    Returns:
-        The output of the command
-    """
-    if desc is not None:
-        print(desc)
-
-    run_kwargs = {
-        "args": command,
-        "shell": True,
-        "env": os.environ if custom_env is None else custom_env,
-        "encoding": 'utf8',
-        "errors": 'ignore'
-    }
-
-    if not live:
-        run_kwargs["stdout"] = run_kwargs["stderr"] = subprocess.PIPE
-
-    result = subprocess.run(check=False, **run_kwargs)
-
-    if result.returncode != 0:
-        error_bits = [
-            f"{error_desc or 'Error running command'}.",
-            f"Command: {command}",
-            f"Error code: {result.returncode}",
-        ]
-        if result.stdout:
-            error_bits.append(f"stdout: {result.stdout}")
-        if result.stderr:
-            error_bits.append(f"stderr: {result.stderr}")
-        raise RuntimeError("\n".join(error_bits))
-
-    return result.stdout or ""
-
-
-# This function was copied from [Fooocus](https://github.com/lllyasviel/Fooocus) repository.
-def run_pip(command, desc=None, live=True):
-    """
-    Run a pip command
-    Args:
-        command: Command to run
-        desc: Description of the command
-        live: Whether to print the output
-    Returns:
-        The output of the command
-    """
-    try:
-        index_url_line = f' --index-url {INDEX_URL}' if INDEX_URL != '' else ''
-        return run_command(
-            command=f'"{PYTHON_EXEC}" -m pip {command} --prefer-binary{index_url_line}',
-            desc=f"Installing {desc}",
-            error_desc=f"Couldn't install {desc}",
-            live=live
-        )
-    except Exception as e:
-        print(f'CMD Failed {command}: {e}')
-        return None
-
-
-def is_installed(package: str) -> bool:
-    """
-    Check if a package is installed
-    Args:
-        package: Package name
-    Returns:
-        Whether the package is installed
-    """
-    try:
-        spec = find_spec(package)
-    except ModuleNotFoundError:
-        return False
-
-    return spec is not None
-
-
-def check_torch_cuda() -> bool:
-    """
-    Check if torch and CUDA is available
-    Returns:
-        Whether CUDA is available
-    """
-    try:
-        import torch
-        return torch.cuda.is_available()
-    except ImportError:
-        return False
-
-
-def requirements_check(requirements_file: str = 'requirements.txt',
-                       pattern: re.Pattern = PATTERN) -> bool:
-    """
-    Check if the requirements file is satisfied
-    Args:
-        requirements_file: Path to the requirements file
-        pattern: Pattern to match the requirements
-    Returns:
-        Whether the requirements file is satisfied
-    """
-    with open(requirements_file, "r", encoding="utf8") as file:
-        for line in file:
-            if line.strip() == "":
-                continue
-
-            m = re.match(pattern, line)
-            if m is None:
-                return False
-
-            package = m.group(1).strip()
-            version_required = (m.group(2) or "").strip()
-
-            if version_required == "":
-                continue
-
-            try:
-                version_installed = metadata.version(package)
-            except Exception:
-                return False
-
-            if version.parse(version_required) != version.parse(version_installed):
-                return False
-
-    return True

fooocusapi/worker.py

@@ -1,1044 +0,0 @@
-"""
-Worker, modify from https://github.com/lllyasviel/Fooocus/blob/main/modules/async_worker.py
-"""
-import copy
-import os
-import random
-import time
-from typing import List
-import logging
-import numpy as np
-import torch
-
-from fooocusapi.models.common.image_meta import image_parse
-from modules.patch import PatchSettings, patch_settings, patch_all
-from modules.flags import Performance
-
-from fooocusapi.utils.file_utils import save_output_file
-from fooocusapi.models.common.task import (
-    GenerationFinishReason,
-    ImageGenerationResult
-)
-from fooocusapi.utils.logger import logger
-from fooocusapi.task_queue import (
-    QueueTask,
-    TaskQueue,
-    TaskOutputs
-)
-
-patch_all()
-
-worker_queue: TaskQueue | None = None
-last_model_name = None
-
-
-def process_stop():
-    """Stop process"""
-    import ldm_patched.modules.model_management
-    ldm_patched.modules.model_management.interrupt_current_processing()
-
-
-@torch.no_grad()
-@torch.inference_mode()
-def task_schedule_loop():
-    """Task schedule loop"""
-    while True:
-        if len(worker_queue.queue) == 0:
-            time.sleep(0.05)
-            continue
-
-        current_task = worker_queue.queue[0]
-        if current_task.start_mills == 0:
-            process_generate(current_task)
-
-
-@torch.no_grad()
-@torch.inference_mode()
-def blocking_get_task_result(job_id: str) -> List[ImageGenerationResult]:
-    """
-    Get task result, when async_task is false
-    :param job_id:
-    :return:
-    """
-    waiting_sleep_steps: int = 0
-    waiting_start_time = time.perf_counter()
-    while not worker_queue.is_task_finished(job_id):
-        if waiting_sleep_steps == 0:
-            logger.std_info(f"[Task Queue] Waiting for task finished, job_id={job_id}")
-        delay = 0.05
-        time.sleep(delay)
-        waiting_sleep_steps += 1
-        if waiting_sleep_steps % int(10 / delay) == 0:
-            waiting_time = time.perf_counter() - waiting_start_time
-            logger.std_info(f"[Task Queue] Already waiting for {round(waiting_time, 1)} seconds, job_id={job_id}")
-
-    task = worker_queue.get_task(job_id, True)
-    return task.task_result
-
-
-@torch.no_grad()
-@torch.inference_mode()
-def process_generate(async_task: QueueTask):
-    """Generate image"""
-    try:
-        import modules.default_pipeline as pipeline
-    except Exception as e:
-        logger.std_error(f'[Task Queue] Import default pipeline error: {e}')
-        if not async_task.is_finished:
-            worker_queue.finish_task(async_task.job_id)
-            async_task.set_result([], True, str(e))
-            logger.std_error(f"[Task Queue] Finish task with error, seq={async_task.job_id}")
-        return []
-
-    import modules.flags as flags
-    import modules.core as core
-    import modules.inpaint_worker as inpaint_worker
-    import modules.config as config
-    import modules.constants as constants
-    import extras.preprocessors as preprocessors
-    import extras.ip_adapter as ip_adapter
-    import extras.face_crop as face_crop
-    import ldm_patched.modules.model_management as model_management
-    from modules.util import (
-        remove_empty_str, HWC3, resize_image,
-        get_image_shape_ceil, set_image_shape_ceil,
-        get_shape_ceil, resample_image, erode_or_dilate,
-        get_enabled_loras, parse_lora_references_from_prompt, apply_wildcards,
-        remove_performance_lora
-    )
-
-    from modules.upscaler import perform_upscale
-    from extras.expansion import safe_str
-    from extras.censor import default_censor
-    from modules.sdxl_styles import (
-        apply_style, get_random_style,
-        fooocus_expansion, apply_arrays, random_style_name
-    )
-
-    pid = os.getpid()
-
-    outputs = TaskOutputs(async_task)
-    results = []
-
-    def refresh_seed(seed_string: int | str | None) -> int:
-        """
-        Refresh and check seed number.
-        :params seed_string: seed, str or int. None means random
-        :return: seed number
-        """
-        if seed_string is None or seed_string == -1:
-            return random.randint(constants.MIN_SEED, constants.MAX_SEED)
-
-        try:
-            seed_value = int(seed_string)
-            if constants.MIN_SEED <= seed_value <= constants.MAX_SEED:
-                return seed_value
-        except ValueError:
-            pass
-        return random.randint(constants.MIN_SEED, constants.MAX_SEED)
-
-    def progressbar(_, number, text):
-        """progress bar"""
-        logger.std_info(f'[Fooocus] {text}')
-        outputs.append(['preview', (number, text, None)])
-
-    def yield_result(_, images, tasks, extension='png',
-                     blockout_nsfw=False, censor=True):
-        """
-        Yield result
-        :param _: async task object
-        :param images: list for generated image
-        :param tasks: the image was generated one by one, when image number is not one, it will be a task list
-        :param extension: extension for saved image
-        :param blockout_nsfw: blockout nsfw image
-        :param censor: censor image
-        :return:
-        """
-        if not isinstance(images, list):
-            images = [images]
-
-        if censor and (config.default_black_out_nsfw or black_out_nsfw):
-            images = default_censor(images)
-
-        results = []
-        for index, im in enumerate(images):
-            if async_task.req_param.save_name == '':
-                image_name = f"{async_task.job_id}-{str(index)}"
-            else:
-                image_name = f"{async_task.req_param.save_name}-{str(index)}"
-            if len(tasks) == 0:
-                img_seed = -1
-                img_meta = {}
-            else:
-                img_seed = tasks[index]['task_seed']
-                img_meta = image_parse(
-                    async_tak=async_task,
-                    task=tasks[index])
-            img_filename = save_output_file(
-                img=im,
-                image_name=image_name,
-                image_meta=img_meta,
-                extension=extension)
-            results.append(ImageGenerationResult(
-                im=img_filename,
-                seed=str(img_seed),
-                finish_reason=GenerationFinishReason.success))
-        async_task.set_result(results, False)
-        worker_queue.finish_task(async_task.job_id)
-        logger.std_info(f"[Task Queue] Finish task, job_id={async_task.job_id}")
-
-        outputs.append(['results', images])
-        pipeline.prepare_text_encoder(async_call=True)
-
-    try:
-        logger.std_info(f"[Task Queue] Task queue start task, job_id={async_task.job_id}")
-        # clear memory
-        global last_model_name
-
-        if last_model_name is None:
-            last_model_name = async_task.req_param.base_model_name
-        if last_model_name != async_task.req_param.base_model_name:
-            model_management.cleanup_models()  # key1
-            model_management.unload_all_models()
-            model_management.soft_empty_cache()  # key2
-            last_model_name = async_task.req_param.base_model_name
-
-        worker_queue.start_task(async_task.job_id)
-
-        execution_start_time = time.perf_counter()
-
-        # Transform parameters
-        params = async_task.req_param
-        prompt = params.prompt
-        negative_prompt = params.negative_prompt
-        style_selections = params.style_selections
-        performance_selection = Performance(params.performance_selection)
-        aspect_ratios_selection = params.aspect_ratios_selection
-        image_number = params.image_number
-        save_metadata_to_images = params.save_meta
-        metadata_scheme = params.meta_scheme
-        save_extension = params.save_extension
-        save_name = params.save_name
-        image_seed = refresh_seed(params.image_seed)
-        read_wildcards_in_order = False
-        sharpness = params.sharpness
-        guidance_scale = params.guidance_scale
-        base_model_name = params.base_model_name
-        refiner_model_name = params.refiner_model_name
-        refiner_switch = params.refiner_switch
-        loras = params.loras
-        input_image_checkbox = params.uov_input_image is not None or params.inpaint_input_image is not None or len(params.image_prompts) > 0
-        current_tab = 'uov' if params.uov_method != flags.disabled else 'ip' if len(params.image_prompts) > 0 else 'inpaint' if params.inpaint_input_image is not None else None
-        uov_method = params.uov_method
-        upscale_value = params.upscale_value
-        uov_input_image = params.uov_input_image
-        outpaint_selections = params.outpaint_selections
-        outpaint_distance_left = params.outpaint_distance_left
-        outpaint_distance_top = params.outpaint_distance_top
-        outpaint_distance_right = params.outpaint_distance_right
-        outpaint_distance_bottom = params.outpaint_distance_bottom
-        inpaint_input_image = params.inpaint_input_image
-        inpaint_additional_prompt = '' if params.inpaint_additional_prompt is None else params.inpaint_additional_prompt
-        inpaint_mask_image_upload = None
-
-        adp = params.advanced_params
-        disable_preview = adp.disable_preview
-        disable_intermediate_results = adp.disable_intermediate_results
-        disable_seed_increment = adp.disable_seed_increment
-        adm_scaler_positive = adp.adm_scaler_positive
-        adm_scaler_negative = adp.adm_scaler_negative
-        adm_scaler_end = adp.adm_scaler_end
-        adaptive_cfg = adp.adaptive_cfg
-        sampler_name = adp.sampler_name
-        scheduler_name = adp.scheduler_name
-        overwrite_step = adp.overwrite_step
-        overwrite_switch = adp.overwrite_switch
-        overwrite_width = adp.overwrite_width
-        overwrite_height = adp.overwrite_height
-        overwrite_vary_strength = adp.overwrite_vary_strength
-        overwrite_upscale_strength = adp.overwrite_upscale_strength
-        mixing_image_prompt_and_vary_upscale = adp.mixing_image_prompt_and_vary_upscale
-        mixing_image_prompt_and_inpaint = adp.mixing_image_prompt_and_inpaint
-        debugging_cn_preprocessor = adp.debugging_cn_preprocessor
-        skipping_cn_preprocessor = adp.skipping_cn_preprocessor
-        canny_low_threshold = adp.canny_low_threshold
-        canny_high_threshold = adp.canny_high_threshold
-        refiner_swap_method = adp.refiner_swap_method
-        controlnet_softness = adp.controlnet_softness
-        freeu_enabled = adp.freeu_enabled
-        freeu_b1 = adp.freeu_b1
-        freeu_b2 = adp.freeu_b2
-        freeu_s1 = adp.freeu_s1
-        freeu_s2 = adp.freeu_s2
-        debugging_inpaint_preprocessor = adp.debugging_inpaint_preprocessor
-        inpaint_disable_initial_latent = adp.inpaint_disable_initial_latent
-        inpaint_engine = adp.inpaint_engine
-        inpaint_strength = adp.inpaint_strength
-        inpaint_respective_field = adp.inpaint_respective_field
-        inpaint_mask_upload_checkbox = adp.inpaint_mask_upload_checkbox
-        invert_mask_checkbox = adp.invert_mask_checkbox
-        inpaint_erode_or_dilate = adp.inpaint_erode_or_dilate
-        black_out_nsfw = adp.black_out_nsfw
-        vae_name = adp.vae_name
-        clip_skip = adp.clip_skip
-
-        cn_tasks = {x: [] for x in flags.ip_list}
-        for img_prompt in params.image_prompts:
-            cn_img, cn_stop, cn_weight, cn_type = img_prompt
-            cn_tasks[cn_type].append([cn_img, cn_stop, cn_weight])
-
-        if inpaint_input_image is not None and inpaint_input_image['image'] is not None:
-            inpaint_image_size = inpaint_input_image['image'].shape[:2]
-            if inpaint_input_image['mask'] is None:
-                inpaint_input_image['mask'] = np.zeros(inpaint_image_size, dtype=np.uint8)
-            else:
-                inpaint_mask_upload_checkbox = True
-
-            inpaint_input_image['mask'] = HWC3(inpaint_input_image['mask'])
-            inpaint_mask_image_upload = inpaint_input_image['mask']
-
-        # Fooocus async_worker.py code start
-
-        outpaint_selections = [o.lower() for o in outpaint_selections]
-        base_model_additional_loras = []
-        raw_style_selections = copy.deepcopy(style_selections)
-        uov_method = uov_method.lower()
-
-        if fooocus_expansion in style_selections:
-            use_expansion = True
-            style_selections.remove(fooocus_expansion)
-        else:
-            use_expansion = False
-
-        use_style = len(style_selections) > 0
-
-        if base_model_name == refiner_model_name:
-            logger.std_warn('[Fooocus] Refiner disabled because base model and refiner are same.')
-            refiner_model_name = 'None'
-
-        steps = performance_selection.steps()
-
-        performance_loras = []
-
-        if performance_selection == Performance.EXTREME_SPEED:
-            logger.std_warn('[Fooocus] Enter LCM mode.')
-            progressbar(async_task, 1, 'Downloading LCM components ...')
-            performance_loras += [(config.downloading_sdxl_lcm_lora(), 1.0)]
-
-            if refiner_model_name != 'None':
-                logger.std_info('[Fooocus] Refiner disabled in LCM mode.')
-
-            refiner_model_name = 'None'
-            sampler_name = 'lcm'
-            scheduler_name = 'lcm'
-            sharpness = 0.0
-            guidance_scale = 1.0
-            adaptive_cfg = 1.0
-            refiner_switch = 1.0
-            adm_scaler_positive = 1.0
-            adm_scaler_negative = 1.0
-            adm_scaler_end = 0.0
-
-        elif performance_selection == Performance.LIGHTNING:
-            logger.std_info('[Fooocus] Enter Lightning mode.')
-            progressbar(async_task, 1, 'Downloading Lightning components ...')
-            performance_loras += [(config.downloading_sdxl_lightning_lora(), 1.0)]
-
-            if refiner_model_name != 'None':
-                logger.std_info('[Fooocus] Refiner disabled in Lightning mode.')
-
-            refiner_model_name = 'None'
-            sampler_name = 'euler'
-            scheduler_name = 'sgm_uniform'
-            sharpness = 0.0
-            guidance_scale = 1.0
-            adaptive_cfg = 1.0
-            refiner_switch = 1.0
-            adm_scaler_positive = 1.0
-            adm_scaler_negative = 1.0
-            adm_scaler_end = 0.0
-
-        elif performance_selection == Performance.HYPER_SD:
-            print('Enter Hyper-SD mode.')
-            progressbar(async_task, 1, 'Downloading Hyper-SD components ...')
-            performance_loras += [(config.downloading_sdxl_hyper_sd_lora(), 0.8)]
-
-            if refiner_model_name != 'None':
-                logger.std_info('[Fooocus] Refiner disabled in Hyper-SD mode.')
-
-            refiner_model_name = 'None'
-            sampler_name = 'dpmpp_sde_gpu'
-            scheduler_name = 'karras'
-            sharpness = 0.0
-            guidance_scale = 1.0
-            adaptive_cfg = 1.0
-            refiner_switch = 1.0
-            adm_scaler_positive = 1.0
-            adm_scaler_negative = 1.0
-            adm_scaler_end = 0.0
-
-        logger.std_info(f'[Parameters] Adaptive CFG = {adaptive_cfg}')
-        logger.std_info(f'[Parameters] CLIP Skip = {clip_skip}')
-        logger.std_info(f'[Parameters] Sharpness = {sharpness}')
-        logger.std_info(f'[Parameters] ControlNet Softness = {controlnet_softness}')
-        logger.std_info(f'[Parameters] ADM Scale = '
-                        f'{adm_scaler_positive} : '
-                        f'{adm_scaler_negative} : '
-                        f'{adm_scaler_end}')
-
-        patch_settings[pid] = PatchSettings(
-            sharpness,
-            adm_scaler_end,
-            adm_scaler_positive,
-            adm_scaler_negative,
-            controlnet_softness,
-            adaptive_cfg
-        )
-
-        cfg_scale = float(guidance_scale)
-        logger.std_info(f'[Parameters] CFG = {cfg_scale}')
-
-        initial_latent = None
-        denoising_strength = 1.0
-        tiled = False
-
-        width, height = aspect_ratios_selection.replace('×', ' ').replace('*', ' ').split(' ')[:2]
-        width, height = int(width), int(height)
-
-        skip_prompt_processing = False
-
-        inpaint_worker.current_task = None
-        inpaint_parameterized = inpaint_engine != 'None'
-        inpaint_image = None
-        inpaint_mask = None
-        inpaint_head_model_path = None
-
-        use_synthetic_refiner = False
-
-        controlnet_canny_path = None
-        controlnet_cpds_path = None
-        clip_vision_path, ip_negative_path, ip_adapter_path, ip_adapter_face_path = None, None, None, None
-
-        seed = int(image_seed)
-        logger.std_info(f'[Parameters] Seed = {seed}')
-
-        goals = []
-        tasks = []
-
-        if input_image_checkbox:
-            if (current_tab == 'uov' or (
-                    current_tab == 'ip' and mixing_image_prompt_and_vary_upscale)) \
-                    and uov_method != flags.disabled and uov_input_image is not None:
-                uov_input_image = HWC3(uov_input_image)
-                if 'vary' in uov_method:
-                    goals.append('vary')
-                elif 'upscale' in uov_method:
-                    goals.append('upscale')
-                    if 'fast' in uov_method:
-                        skip_prompt_processing = True
-                    else:
-                        steps = performance_selection.steps_uov()
-
-                    progressbar(async_task, 1, 'Downloading upscale models ...')
-                    config.downloading_upscale_model()
-            if (current_tab == 'inpaint' or (
-                    current_tab == 'ip' and mixing_image_prompt_and_inpaint)) \
-                    and isinstance(inpaint_input_image, dict):
-                inpaint_image = inpaint_input_image['image']
-                inpaint_mask = inpaint_input_image['mask'][:, :, 0]
-
-                if inpaint_mask_upload_checkbox:
-                    if isinstance(inpaint_mask_image_upload, np.ndarray):
-                        if inpaint_mask_image_upload.ndim == 3:
-                            H, W, C = inpaint_image.shape
-                            inpaint_mask_image_upload = resample_image(inpaint_mask_image_upload, width=W, height=H)
-                            inpaint_mask_image_upload = np.mean(inpaint_mask_image_upload, axis=2)
-                            inpaint_mask_image_upload = (inpaint_mask_image_upload > 127).astype(np.uint8) * 255
-                            inpaint_mask = np.maximum(np.zeros(shape=(H, W), dtype=np.uint8), inpaint_mask_image_upload)
-
-                if int(inpaint_erode_or_dilate) != 0:
-                    inpaint_mask = erode_or_dilate(inpaint_mask, inpaint_erode_or_dilate)
-
-                if invert_mask_checkbox:
-                    inpaint_mask = 255 - inpaint_mask
-
-                inpaint_image = HWC3(inpaint_image)
-                if isinstance(inpaint_image, np.ndarray) and isinstance(inpaint_mask, np.ndarray) \
-                        and (np.any(inpaint_mask > 127) or len(outpaint_selections) > 0):
-                    progressbar(async_task, 1, 'Downloading upscale models ...')
-                    config.downloading_upscale_model()
-                    if inpaint_parameterized:
-                        progressbar(async_task, 1, 'Downloading inpainter ...')
-                        inpaint_head_model_path, inpaint_patch_model_path = config.downloading_inpaint_models(
-                            inpaint_engine)
-                        base_model_additional_loras += [(inpaint_patch_model_path, 1.0)]
-                        logger.std_info(f'[Inpaint] Current inpaint model is {inpaint_patch_model_path}')
-                        if refiner_model_name == 'None':
-                            use_synthetic_refiner = True
-                            refiner_switch = 0.8
-                    else:
-                        inpaint_head_model_path, inpaint_patch_model_path = None, None
-                        logger.std_info('[Inpaint] Parameterized inpaint is disabled.')
-                    if inpaint_additional_prompt != '':
-                        if prompt == '':
-                            prompt = inpaint_additional_prompt
-                        else:
-                            prompt = inpaint_additional_prompt + '\n' + prompt
-                    goals.append('inpaint')
-            if current_tab == 'ip' or \
-                    mixing_image_prompt_and_vary_upscale or \
-                    mixing_image_prompt_and_inpaint:
-                goals.append('cn')
-                progressbar(async_task, 1, 'Downloading control models ...')
-                if len(cn_tasks[flags.cn_canny]) > 0:
-                    controlnet_canny_path = config.downloading_controlnet_canny()
-                if len(cn_tasks[flags.cn_cpds]) > 0:
-                    controlnet_cpds_path = config.downloading_controlnet_cpds()
-                if len(cn_tasks[flags.cn_ip]) > 0:
-                    clip_vision_path, ip_negative_path, ip_adapter_path = config.downloading_ip_adapters('ip')
-                if len(cn_tasks[flags.cn_ip_face]) > 0:
-                    clip_vision_path, ip_negative_path, ip_adapter_face_path = config.downloading_ip_adapters(
-                        'face')
-                progressbar(async_task, 1, 'Loading control models ...')
-
-        # Load or unload CNs
-        pipeline.refresh_controlnets([controlnet_canny_path, controlnet_cpds_path])
-        ip_adapter.load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_path)
-        ip_adapter.load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_face_path)
-
-        if overwrite_step > 0:
-            steps = overwrite_step
-
-        switch = int(round(steps * refiner_switch))
-
-        if overwrite_switch > 0:
-            switch = overwrite_switch
-
-        if overwrite_width > 0:
-            width = overwrite_width
-
-        if overwrite_height > 0:
-            height = overwrite_height
-
-        logger.std_info(f'[Parameters] Sampler = {sampler_name} - {scheduler_name}')
-        logger.std_info(f'[Parameters] Steps = {steps} - {switch}')
-
-        progressbar(async_task, 1, 'Initializing ...')
-
-        if not skip_prompt_processing:
-
-            prompts = remove_empty_str([safe_str(p) for p in prompt.splitlines()], default='')
-            negative_prompts = remove_empty_str([safe_str(p) for p in negative_prompt.splitlines()], default='')
-
-            prompt = prompts[0]
-            negative_prompt = negative_prompts[0]
-
-            if prompt == '':
-                # disable expansion when empty since it is not meaningful and influences image prompt
-                use_expansion = False
-
-            extra_positive_prompts = prompts[1:] if len(prompts) > 1 else []
-            extra_negative_prompts = negative_prompts[1:] if len(negative_prompts) > 1 else []
-
-            progressbar(async_task, 3, 'Loading models ...')
-            lora_filenames = remove_performance_lora(config.lora_filenames, performance_selection)
-            loras, prompt = parse_lora_references_from_prompt(prompt, loras, config.default_max_lora_number, lora_filenames=lora_filenames)
-            loras += performance_loras
-
-            pipeline.refresh_everything(
-                refiner_model_name=refiner_model_name,
-                base_model_name=base_model_name,
-                loras=loras,
-                base_model_additional_loras=base_model_additional_loras,
-                use_synthetic_refiner=use_synthetic_refiner)
-
-            pipeline.set_clip_skip(clip_skip)
-
-            progressbar(async_task, 3, 'Processing prompts ...')
-            tasks = []
-
-            for i in range(image_number):
-                if disable_seed_increment:
-                    task_seed = seed % (constants.MAX_SEED + 1)
-                else:
-                    task_seed = (seed + i) % (constants.MAX_SEED + 1)  # randint is inclusive, % is not
-
-                task_rng = random.Random(task_seed)  # may bind to inpaint noise in the future
-                task_prompt = apply_wildcards(prompt, task_rng, i, read_wildcards_in_order)
-                task_prompt = apply_arrays(task_prompt, i)
-                task_negative_prompt = apply_wildcards(negative_prompt, task_rng, i, read_wildcards_in_order)
-                task_extra_positive_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in
-                                               extra_positive_prompts]
-                task_extra_negative_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in
-                                               extra_negative_prompts]
-
-                positive_basic_workloads = []
-                negative_basic_workloads = []
-
-                task_styles = style_selections.copy()
-                if use_style:
-                    for index, style in enumerate(task_styles):
-                        if style == random_style_name:
-                            style = get_random_style(task_rng)
-                            task_styles[index] = style
-                        p, n = apply_style(style, positive=task_prompt)
-                        positive_basic_workloads = positive_basic_workloads + p
-                        negative_basic_workloads = negative_basic_workloads + n
-                else:
-                    positive_basic_workloads.append(task_prompt)
-
-                negative_basic_workloads.append(task_negative_prompt)  # Always use independent workload for negative.
-
-                positive_basic_workloads = positive_basic_workloads + task_extra_positive_prompts
-                negative_basic_workloads = negative_basic_workloads + task_extra_negative_prompts
-
-                positive_basic_workloads = remove_empty_str(positive_basic_workloads, default=task_prompt)
-                negative_basic_workloads = remove_empty_str(negative_basic_workloads, default=task_negative_prompt)
-
-                tasks.append(dict(
-                    task_seed=task_seed,
-                    task_prompt=task_prompt,
-                    task_negative_prompt=task_negative_prompt,
-                    positive=positive_basic_workloads,
-                    negative=negative_basic_workloads,
-                    expansion='',
-                    c=None,
-                    uc=None,
-                    positive_top_k=len(positive_basic_workloads),
-                    negative_top_k=len(negative_basic_workloads),
-                    log_positive_prompt='\n'.join([task_prompt] + task_extra_positive_prompts),
-                    log_negative_prompt='\n'.join([task_negative_prompt] + task_extra_negative_prompts),
-                    styles=task_styles
-                ))
-
-            if use_expansion:
-                for i, t in enumerate(tasks):
-                    progressbar(async_task, 4, f'Preparing Fooocus text #{i + 1} ...')
-                    expansion = pipeline.final_expansion(t['task_prompt'], t['task_seed'])
-                    logger.std_info(f'[Prompt Expansion] {expansion}')
-                    t['expansion'] = expansion
-                    t['positive'] = copy.deepcopy(t['positive']) + [expansion]  # Deep copy.
-
-            for i, t in enumerate(tasks):
-                progressbar(async_task, 5, f'Encoding positive #{i + 1} ...')
-                t['c'] = pipeline.clip_encode(texts=t['positive'], pool_top_k=t['positive_top_k'])
-
-            for i, t in enumerate(tasks):
-                if abs(float(cfg_scale) - 1.0) < 1e-4:
-                    t['uc'] = pipeline.clone_cond(t['c'])
-                else:
-                    progressbar(async_task, 6, f'Encoding negative #{i + 1} ...')
-                    t['uc'] = pipeline.clip_encode(texts=t['negative'], pool_top_k=t['negative_top_k'])
-
-        if len(goals) > 0:
-            progressbar(async_task, 7, 'Image processing ...')
-
-        if 'vary' in goals:
-            if 'subtle' in uov_method:
-                denoising_strength = 0.5
-            if 'strong' in uov_method:
-                denoising_strength = 0.85
-            if overwrite_vary_strength > 0:
-                denoising_strength = overwrite_vary_strength
-
-            shape_ceil = get_image_shape_ceil(uov_input_image)
-            if shape_ceil < 1024:
-                logger.std_warn('[Vary] Image is resized because it is too small.')
-                shape_ceil = 1024
-            elif shape_ceil > 2048:
-                logger.std_warn('[Vary] Image is resized because it is too big.')
-                shape_ceil = 2048
-
-            uov_input_image = set_image_shape_ceil(uov_input_image, shape_ceil)
-
-            initial_pixels = core.numpy_to_pytorch(uov_input_image)
-            progressbar(async_task, 8, 'VAE encoding ...')
-
-            candidate_vae, _ = pipeline.get_candidate_vae(
-                steps=steps,
-                switch=switch,
-                denoise=denoising_strength,
-                refiner_swap_method=refiner_swap_method
-            )
-
-            initial_latent = core.encode_vae(vae=candidate_vae, pixels=initial_pixels)
-            B, C, H, W = initial_latent['samples'].shape
-            width = W * 8
-            height = H * 8
-            logger.std_info(f'[Vary] Final resolution is {str((height, width))}.')
-
-        if 'upscale' in goals:
-            H, W, C = uov_input_image.shape
-            progressbar(async_task, 9, f'Upscaling image from {str((H, W))} ...')
-            uov_input_image = perform_upscale(uov_input_image)
-            logger.std_info('[Upscale] Image upscale.')
-
-            if upscale_value is not None and upscale_value > 1.0:
-                f = upscale_value
-            else:
-                if '1.5x' in uov_method:
-                    f = 1.5
-                elif '2x' in uov_method:
-                    f = 2.0
-                else:
-                    f = 1.0
-
-            shape_ceil = get_shape_ceil(H * f, W * f)
-
-            if shape_ceil < 1024:
-                logger.std_info('[Upscale] Image is resized because it is too small.')
-                uov_input_image = set_image_shape_ceil(uov_input_image, 1024)
-                shape_ceil = 1024
-            else:
-                uov_input_image = resample_image(uov_input_image, width=W * f, height=H * f)
-
-            image_is_super_large = shape_ceil > 2800
-
-            if 'fast' in uov_method:
-                direct_return = True
-            elif image_is_super_large:
-                logger.std_info('[Upscale] Image is too large. Directly returned the SR image. '
-                                'Usually directly return SR image at 4K resolution '
-                                'yields better results than SDXL diffusion.')
-                direct_return = True
-            else:
-                direct_return = False
-
-            if direct_return:
-                # d = [('Upscale (Fast)', '2x')]
-                # log(uov_input_image, d, output_format=save_extension)
-                if config.default_black_out_nsfw or black_out_nsfw:
-                    uov_input_image = default_censor(uov_input_image)
-                yield_result(async_task, uov_input_image, tasks, save_extension, False, False)
-                return
-
-            tiled = True
-            denoising_strength = 0.382
-
-            if overwrite_upscale_strength > 0:
-                denoising_strength = overwrite_upscale_strength
-
-            initial_pixels = core.numpy_to_pytorch(uov_input_image)
-            progressbar(async_task, 10, 'VAE encoding ...')
-
-            candidate_vae, _ = pipeline.get_candidate_vae(
-                steps=steps,
-                switch=switch,
-                denoise=denoising_strength,
-                refiner_swap_method=refiner_swap_method
-            )
-
-            initial_latent = core.encode_vae(
-                vae=candidate_vae,
-                pixels=initial_pixels, tiled=True)
-            B, C, H, W = initial_latent['samples'].shape
-            width = W * 8
-            height = H * 8
-            logger.std_info(f'[Upscale] Final resolution is {str((height, width))}.')
-
-        if 'inpaint' in goals:
-            if len(outpaint_selections) > 0:
-                H, W, C = inpaint_image.shape
-                if 'top' in outpaint_selections:
-                    distance_top = int(H * 0.3)
-                    if outpaint_distance_top > 0:
-                        distance_top = outpaint_distance_top
-
-                    inpaint_image = np.pad(inpaint_image, [[distance_top, 0], [0, 0], [0, 0]], mode='edge')
-                    inpaint_mask = np.pad(inpaint_mask, [[distance_top, 0], [0, 0]], mode='constant',
-                                          constant_values=255)
-
-                if 'bottom' in outpaint_selections:
-                    distance_bottom = int(H * 0.3)
-                    if outpaint_distance_bottom > 0:
-                        distance_bottom = outpaint_distance_bottom
-
-                    inpaint_image = np.pad(inpaint_image, [[0, distance_bottom], [0, 0], [0, 0]], mode='edge')
-                    inpaint_mask = np.pad(inpaint_mask, [[0, distance_bottom], [0, 0]], mode='constant',
-                                          constant_values=255)
-
-                H, W, C = inpaint_image.shape
-                if 'left' in outpaint_selections:
-                    distance_left = int(W * 0.3)
-                    if outpaint_distance_left > 0:
-                        distance_left = outpaint_distance_left
-
-                    inpaint_image = np.pad(inpaint_image, [[0, 0], [distance_left, 0], [0, 0]], mode='edge')
-                    inpaint_mask = np.pad(inpaint_mask, [[0, 0], [distance_left, 0]], mode='constant',
-                                          constant_values=255)
-
-                if 'right' in outpaint_selections:
-                    distance_right = int(W * 0.3)
-                    if outpaint_distance_right > 0:
-                        distance_right = outpaint_distance_right
-
-                    inpaint_image = np.pad(inpaint_image, [[0, 0], [0, distance_right], [0, 0]], mode='edge')
-                    inpaint_mask = np.pad(inpaint_mask, [[0, 0], [0, distance_right]], mode='constant',
-                                          constant_values=255)
-
-                inpaint_image = np.ascontiguousarray(inpaint_image.copy())
-                inpaint_mask = np.ascontiguousarray(inpaint_mask.copy())
-                inpaint_strength = 1.0
-                inpaint_respective_field = 1.0
-
-            denoising_strength = inpaint_strength
-
-            inpaint_worker.current_task = inpaint_worker.InpaintWorker(
-                image=inpaint_image,
-                mask=inpaint_mask,
-                use_fill=denoising_strength > 0.99,
-                k=inpaint_respective_field
-            )
-
-            if debugging_inpaint_preprocessor:
-                yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), tasks,
-                             black_out_nsfw)
-                return
-
-            progressbar(async_task, 11, 'VAE Inpaint encoding ...')
-
-            inpaint_pixel_fill = core.numpy_to_pytorch(inpaint_worker.current_task.interested_fill)
-            inpaint_pixel_image = core.numpy_to_pytorch(inpaint_worker.current_task.interested_image)
-            inpaint_pixel_mask = core.numpy_to_pytorch(inpaint_worker.current_task.interested_mask)
-
-            candidate_vae, candidate_vae_swap = pipeline.get_candidate_vae(
-                steps=steps,
-                switch=switch,
-                denoise=denoising_strength,
-                refiner_swap_method=refiner_swap_method
-            )
-
-            latent_inpaint, latent_mask = core.encode_vae_inpaint(
-                mask=inpaint_pixel_mask,
-                vae=candidate_vae,
-                pixels=inpaint_pixel_image)
-
-            latent_swap = None
-            if candidate_vae_swap is not None:
-                progressbar(async_task, 12, 'VAE SD15 encoding ...')
-                latent_swap = core.encode_vae(
-                    vae=candidate_vae_swap,
-                    pixels=inpaint_pixel_fill)['samples']
-
-            progressbar(async_task, 13, 'VAE encoding ...')
-            latent_fill = core.encode_vae(
-                vae=candidate_vae,
-                pixels=inpaint_pixel_fill)['samples']
-
-            inpaint_worker.current_task.load_latent(
-                latent_fill=latent_fill, latent_mask=latent_mask, latent_swap=latent_swap)
-
-            if inpaint_parameterized:
-                pipeline.final_unet = inpaint_worker.current_task.patch(
-                    inpaint_head_model_path=inpaint_head_model_path,
-                    inpaint_latent=latent_inpaint,
-                    inpaint_latent_mask=latent_mask,
-                    model=pipeline.final_unet
-                )
-
-            if not inpaint_disable_initial_latent:
-                initial_latent = {'samples': latent_fill}
-
-            B, C, H, W = latent_fill.shape
-            height, width = H * 8, W * 8
-            final_height, final_width = inpaint_worker.current_task.image.shape[:2]
-            logger.std_info(f'[Inpaint] Final resolution is {str((final_height, final_width))}, latent is {str((height, width))}.')
-
-        if 'cn' in goals:
-            for task in cn_tasks[flags.cn_canny]:
-                cn_img, cn_stop, cn_weight = task
-                cn_img = resize_image(HWC3(cn_img), width=width, height=height)
-
-                if not skipping_cn_preprocessor:
-                    cn_img = preprocessors.canny_pyramid(cn_img, canny_low_threshold, canny_high_threshold)
-
-                cn_img = HWC3(cn_img)
-                task[0] = core.numpy_to_pytorch(cn_img)
-                if debugging_cn_preprocessor:
-                    yield_result(async_task, cn_img, tasks, save_extension, black_out_nsfw)
-                    return
-            for task in cn_tasks[flags.cn_cpds]:
-                cn_img, cn_stop, cn_weight = task
-                cn_img = resize_image(HWC3(cn_img), width=width, height=height)
-
-                if not skipping_cn_preprocessor:
-                    cn_img = preprocessors.cpds(cn_img)
-
-                cn_img = HWC3(cn_img)
-                task[0] = core.numpy_to_pytorch(cn_img)
-                if debugging_cn_preprocessor:
-                    yield_result(async_task, cn_img, tasks, save_extension, black_out_nsfw)
-                    return
-            for task in cn_tasks[flags.cn_ip]:
-                cn_img, cn_stop, cn_weight = task
-                cn_img = HWC3(cn_img)
-
-                # https://github.com/tencent-ailab/IP-Adapter/blob/d580c50a291566bbf9fc7ac0f760506607297e6d/README.md?plain=1#L75
-                cn_img = resize_image(cn_img, width=224, height=224, resize_mode=0)
-
-                task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_path)
-                if debugging_cn_preprocessor:
-                    yield_result(async_task, cn_img, tasks, save_extension, black_out_nsfw)
-                    return
-            for task in cn_tasks[flags.cn_ip_face]:
-                cn_img, cn_stop, cn_weight = task
-                cn_img = HWC3(cn_img)
-
-                if not skipping_cn_preprocessor:
-                    cn_img = face_crop.crop_image(cn_img)
-
-                # https://github.com/tencent-ailab/IP-Adapter/blob/d580c50a291566bbf9fc7ac0f760506607297e6d/README.md?plain=1#L75
-                cn_img = resize_image(cn_img, width=224, height=224, resize_mode=0)
-
-                task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_face_path)
-                if debugging_cn_preprocessor:
-                    yield_result(async_task, cn_img, tasks, save_extension, black_out_nsfw)
-                    return
-
-            all_ip_tasks = cn_tasks[flags.cn_ip] + cn_tasks[flags.cn_ip_face]
-
-            if len(all_ip_tasks) > 0:
-                pipeline.final_unet = ip_adapter.patch_model(pipeline.final_unet, all_ip_tasks)
-
-        if freeu_enabled:
-            logger.std_info('[Fooocus] FreeU is enabled!')
-            pipeline.final_unet = core.apply_freeu(
-                pipeline.final_unet,
-                freeu_b1,
-                freeu_b2,
-                freeu_s1,
-                freeu_s2
-            )
-
-        all_steps = steps * image_number
-
-        logger.std_info(f'[Parameters] Denoising Strength = {denoising_strength}')
-
-        if isinstance(initial_latent, dict) and 'samples' in initial_latent:
-            log_shape = initial_latent['samples'].shape
-        else:
-            log_shape = f'Image Space {(height, width)}'
-
-        logger.std_info(f'[Parameters] Initial Latent shape: {log_shape}')
-
-        preparation_time = time.perf_counter() - execution_start_time
-        logger.std_info(f'[Fooocus] Preparation time: {preparation_time:.2f} seconds')
-
-        final_sampler_name = sampler_name
-        final_scheduler_name = scheduler_name
-
-        if scheduler_name in ['lcm', 'tcd']:
-            final_scheduler_name = 'sgm_uniform'
-
-            def patch_discrete(unet):
-                return core.opModelSamplingDiscrete.patch(
-                    pipeline.final_unet,
-                    sampling=scheduler_name,
-                    zsnr=False)[0]
-
-            if pipeline.final_unet is not None:
-                pipeline.final_unet = patch_discrete(pipeline.final_unet)
-            if pipeline.final_refiner_unet is not None:
-                pipeline.final_refiner_unet = patch_discrete(pipeline.final_refiner_unet)
-            logger.std_info(f'[Fooocus] Using {scheduler_name} scheduler.')
-        elif scheduler_name == 'edm_playground_v2.5':
-            final_scheduler_name = 'karras'
-
-            def patch_edm(unet):
-                return core.opModelSamplingContinuousEDM.patch(
-                    unet,
-                    sampling=scheduler_name,
-                    sigma_max=120.0,
-                    sigma_min=0.002)[0]
-
-            if pipeline.final_unet is not None:
-                pipeline.final_unet = patch_edm(pipeline.final_unet)
-            if pipeline.final_refiner_unet is not None:
-                pipeline.final_refiner_unet = patch_edm(pipeline.final_refiner_unet)
-
-            logger.std_info(f'[Fooocus] Using {scheduler_name} scheduler.')
-
-        outputs.append(['preview', (13, 'Moving model to GPU ...', None)])
-
-        def callback(step, x0, x, total_steps, y):
-            """callback, used for progress and preview"""
-            done_steps = current_task_id * steps + step
-            outputs.append(['preview', (
-                int(15.0 + 85.0 * float(done_steps) / float(all_steps)),
-                f'Step {step}/{total_steps} in the {current_task_id + 1}-th Sampling',
-                y)])
-
-        for current_task_id, task in enumerate(tasks):
-            execution_start_time = time.perf_counter()
-
-            try:
-                positive_cond, negative_cond = task['c'], task['uc']
-
-                if 'cn' in goals:
-                    for cn_flag, cn_path in [
-                        (flags.cn_canny, controlnet_canny_path),
-                        (flags.cn_cpds, controlnet_cpds_path)
-                    ]:
-                        for cn_img, cn_stop, cn_weight in cn_tasks[cn_flag]:
-                            positive_cond, negative_cond = core.apply_controlnet(
-                                positive_cond, negative_cond,
-                                pipeline.loaded_ControlNets[cn_path], cn_img, cn_weight, 0, cn_stop)
-
-                imgs = pipeline.process_diffusion(
-                    positive_cond=positive_cond,
-                    negative_cond=negative_cond,
-                    steps=steps,
-                    switch=switch,
-                    width=width,
-                    height=height,
-                    image_seed=task['task_seed'],
-                    callback=callback,
-                    sampler_name=final_sampler_name,
-                    scheduler_name=final_scheduler_name,
-                    latent=initial_latent,
-                    denoise=denoising_strength,
-                    tiled=tiled,
-                    cfg_scale=cfg_scale,
-                    refiner_swap_method=refiner_swap_method,
-                    disable_preview=disable_preview
-                )
-
-                del task['c'], task['uc'], positive_cond, negative_cond  # Save memory
-
-                if inpaint_worker.current_task is not None:
-                    imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
-
-                # Fooocus async_worker.py code end
-
-                results += imgs
-            except model_management.InterruptProcessingException as e:
-                logger.std_warn("[Fooocus] User stopped")
-                results = []
-                results.append(ImageGenerationResult(
-                    im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.user_cancel))
-                async_task.set_result(results, True, str(e))
-                break
-            except Exception as e:
-                logger.std_error(f'[Fooocus] Process error: {e}')
-                logging.exception(e)
-                results = []
-                results.append(ImageGenerationResult(
-                    im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.error))
-                async_task.set_result(results, True, str(e))
-                break
-
-            execution_time = time.perf_counter() - execution_start_time
-            logger.std_info(f'[Fooocus] Generating and saving time: {execution_time:.2f} seconds')
-
-        if async_task.finish_with_error:
-            worker_queue.finish_task(async_task.job_id)
-            return async_task.task_result
-        yield_result(None, results, tasks, save_extension, black_out_nsfw)
-        return
-    except Exception as e:
-        logger.std_error(f'[Fooocus] Worker error: {e}')
-
-        if not async_task.is_finished:
-            async_task.set_result([], True, str(e))
-            worker_queue.finish_task(async_task.job_id)
-            logger.std_info(f"[Task Queue] Finish task with error, job_id={async_task.job_id}")

predict.py

@@ -1,316 +0,0 @@
-"""
-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)

repositories/Fooocus/__init__.py

@@ -1,4 +0,0 @@
-"""
-Created By: ishwor subedi
-Date: 2024-07-19
-"""

repositories/Fooocus/args_manager.py

@@ -1,55 +0,0 @@
-import ldm_patched.modules.args_parser as args_parser
-
-args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.")
-
-args_parser.parser.add_argument("--preset", type=str, default=None, help="Apply specified UI preset.")
-args_parser.parser.add_argument("--disable-preset-selection", action='store_true',
-                                help="Disables preset selection in Gradio.")
-
-args_parser.parser.add_argument("--language", type=str, default='default',
-                                help="Translate UI using json files in [language] folder. "
-                                  "For example, [--language example] will use [language/example.json] for translation.")
-
-# For example, https://github.com/lllyasviel/Fooocus/issues/849
-args_parser.parser.add_argument("--disable-offload-from-vram", action="store_true",
-                                help="Force loading models to vram when the unload can be avoided. "
-                                  "Some Mac users may need this.")
-
-args_parser.parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)
-args_parser.parser.add_argument("--disable-image-log", action='store_true',
-                                help="Prevent writing images and logs to hard drive.")
-
-args_parser.parser.add_argument("--disable-analytics", action='store_true',
-                                help="Disables analytics for Gradio.")
-
-args_parser.parser.add_argument("--disable-metadata", action='store_true',
-                                help="Disables saving metadata to images.")
-
-args_parser.parser.add_argument("--disable-preset-download", action='store_true',
-                                help="Disables downloading models for presets", default=False)
-
-args_parser.parser.add_argument("--enable-describe-uov-image", action='store_true',
-                                help="Disables automatic description of uov images when prompt is empty", default=False)
-
-args_parser.parser.add_argument("--always-download-new-model", action='store_true',
-                                help="Always download newer models ", default=False)
-
-args_parser.parser.set_defaults(
-    disable_cuda_malloc=True,
-    in_browser=True,
-    port=None
-)
-
-args_parser.args = args_parser.parser.parse_args()
-
-# (Disable by default because of issues like https://github.com/lllyasviel/Fooocus/issues/724)
-args_parser.args.always_offload_from_vram = not args_parser.args.disable_offload_from_vram
-
-if args_parser.args.disable_analytics:
-    import os
-    os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
-
-if args_parser.args.disable_in_browser:
-    args_parser.args.in_browser = False
-
-args = args_parser.args

repositories/Fooocus/extras/BLIP/configs/bert_config.json

@@ -1,21 +0,0 @@
-{
-  "architectures": [
-    "BertModel"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 768,
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 512,
-  "model_type": "bert",
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
-  "pad_token_id": 0,
-  "type_vocab_size": 2,
-  "vocab_size": 30522,
-  "encoder_width": 768,
-  "add_cross_attention": true   
-}

repositories/Fooocus/extras/BLIP/configs/caption_coco.yaml

@@ -1,33 +0,0 @@
-image_root: '/export/share/datasets/vision/coco/images/'
-ann_root: 'annotation'
-coco_gt_root: 'annotation/coco_gt'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
-
-# size of vit model; base or large
-vit: 'base'
-vit_grad_ckpt: False
-vit_ckpt_layer: 0
-batch_size: 32
-init_lr: 1e-5
-
-# vit: 'large'
-# vit_grad_ckpt: True
-# vit_ckpt_layer: 5
-# batch_size: 16
-# init_lr: 2e-6
-
-image_size: 384
-
-# generation configs
-max_length: 20  
-min_length: 5
-num_beams: 3
-prompt: 'a picture of '
-
-# optimizer
-weight_decay: 0.05
-min_lr: 0
-max_epoch: 5
-

repositories/Fooocus/extras/BLIP/configs/med_config.json

@@ -1,21 +0,0 @@
-{
-  "architectures": [
-    "BertModel"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 768,
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 512,
-  "model_type": "bert",
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
-  "pad_token_id": 0,
-  "type_vocab_size": 2,
-  "vocab_size": 30524,
-  "encoder_width": 768,
-  "add_cross_attention": true   
-}

repositories/Fooocus/extras/BLIP/configs/nlvr.yaml

@@ -1,21 +0,0 @@
-image_root: '/export/share/datasets/vision/NLVR2/' 
-ann_root: 'annotation'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_nlvr.pth'
-
-#size of vit model; base or large
-vit: 'base'
-batch_size_train: 16 
-batch_size_test: 64 
-vit_grad_ckpt: False
-vit_ckpt_layer: 0
-max_epoch: 15
-
-image_size: 384
-
-# optimizer
-weight_decay: 0.05
-init_lr: 3e-5
-min_lr: 0
-

repositories/Fooocus/extras/BLIP/configs/nocaps.yaml

@@ -1,15 +0,0 @@
-image_root: '/export/share/datasets/vision/nocaps/'
-ann_root: 'annotation'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
-
-vit: 'base'
-batch_size: 32
-
-image_size: 384
-
-max_length: 20
-min_length: 5
-num_beams: 3
-prompt: 'a picture of '

repositories/Fooocus/extras/BLIP/configs/pretrain.yaml

@@ -1,27 +0,0 @@
-train_file: ['/export/share/junnan-li/VL_pretrain/annotation/coco_karpathy_train.json',
-             '/export/share/junnan-li/VL_pretrain/annotation/vg_caption.json',
-             ]
-laion_path: ''   
-
-# size of vit model; base or large
-vit: 'base'
-vit_grad_ckpt: False
-vit_ckpt_layer: 0
-
-image_size: 224
-batch_size: 75
-
-queue_size: 57600
-alpha: 0.4
-
-# optimizer
-weight_decay: 0.05
-init_lr: 3e-4
-min_lr: 1e-6
-warmup_lr: 1e-6
-lr_decay_rate: 0.9
-max_epoch: 20
-warmup_steps: 3000
-
-
-

repositories/Fooocus/extras/BLIP/configs/retrieval_coco.yaml

@@ -1,34 +0,0 @@
-image_root: '/export/share/datasets/vision/coco/images/'
-ann_root: 'annotation'
-dataset: 'coco'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'
-
-# size of vit model; base or large
-
-vit: 'base'
-batch_size_train: 32
-batch_size_test: 64
-vit_grad_ckpt: True
-vit_ckpt_layer: 4
-init_lr: 1e-5
-
-# vit: 'large'
-# batch_size_train: 16
-# batch_size_test: 32
-# vit_grad_ckpt: True
-# vit_ckpt_layer: 12
-# init_lr: 5e-6
-
-image_size: 384
-queue_size: 57600
-alpha: 0.4
-k_test: 256
-negative_all_rank: True
-
-# optimizer
-weight_decay: 0.05
-min_lr: 0
-max_epoch: 6
-

repositories/Fooocus/extras/BLIP/configs/retrieval_flickr.yaml

@@ -1,34 +0,0 @@
-image_root: '/export/share/datasets/vision/flickr30k/'
-ann_root: 'annotation'
-dataset: 'flickr'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_flickr.pth'
-
-# size of vit model; base or large
-
-vit: 'base'
-batch_size_train: 32
-batch_size_test: 64
-vit_grad_ckpt: True
-vit_ckpt_layer: 4
-init_lr: 1e-5
-
-# vit: 'large'
-# batch_size_train: 16
-# batch_size_test: 32
-# vit_grad_ckpt: True
-# vit_ckpt_layer: 10
-# init_lr: 5e-6
-
-image_size: 384
-queue_size: 57600
-alpha: 0.4
-k_test: 128
-negative_all_rank: False
-
-# optimizer
-weight_decay: 0.05
-min_lr: 0
-max_epoch: 6
-

repositories/Fooocus/extras/BLIP/configs/retrieval_msrvtt.yaml

@@ -1,12 +0,0 @@
-video_root: '/export/share/dongxuli/data/msrvtt_retrieval/videos'
-ann_root: 'annotation'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'
-
-# size of vit model; base or large
-vit: 'base'
-batch_size: 64
-k_test: 128
-image_size: 384
-num_frm_test: 8

repositories/Fooocus/extras/BLIP/configs/vqa.yaml

@@ -1,25 +0,0 @@
-vqa_root: '/export/share/datasets/vision/VQA/Images/mscoco/' #followed by train2014/
-vg_root: '/export/share/datasets/vision/visual-genome/'  #followed by image/
-train_files: ['vqa_train','vqa_val','vg_qa']
-ann_root: 'annotation'
-
-# set pretrained as a file path or an url
-pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'
-
-# size of vit model; base or large
-vit: 'base'
-batch_size_train: 16 
-batch_size_test: 32 
-vit_grad_ckpt: False
-vit_ckpt_layer: 0
-init_lr: 2e-5
-
-image_size: 480
-
-k_test: 128
-inference: 'rank'
-
-# optimizer
-weight_decay: 0.05
-min_lr: 0
-max_epoch: 10

repositories/Fooocus/extras/BLIP/models/bert_tokenizer/config.json

@@ -1,23 +0,0 @@
-{
-  "architectures": [
-    "BertForMaskedLM"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "gradient_checkpointing": false,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 768,
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 512,
-  "model_type": "bert",
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
-  "pad_token_id": 0,
-  "position_embedding_type": "absolute",
-  "transformers_version": "4.6.0.dev0",
-  "type_vocab_size": 2,
-  "use_cache": true,
-  "vocab_size": 30522
-}

repositories/Fooocus/extras/BLIP/models/bert_tokenizer/tokenizer_config.json

@@ -1,3 +0,0 @@
-{
-  "do_lower_case": true
-}

repositories/Fooocus/extras/BLIP/models/blip.py

@@ -1,239 +0,0 @@
-'''
- * Copyright (c) 2022, salesforce.com, inc.
- * All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
- * By Junnan Li
-'''
-import warnings
-warnings.filterwarnings("ignore")
-
-from extras.BLIP.models.vit import VisionTransformer, interpolate_pos_embed
-from extras.BLIP.models.med import BertConfig, BertModel, BertLMHeadModel
-from transformers import BertTokenizer
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-import os
-from urllib.parse import urlparse
-from timm.models.hub import download_cached_file
-
-class BLIP_Base(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/med_config.json',  
-                 image_size = 224,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,                 
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """               
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
-        self.tokenizer = init_tokenizer()   
-        med_config = BertConfig.from_json_file(med_config)
-        med_config.encoder_width = vision_width
-        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)  
-
-        
-    def forward(self, image, caption, mode):
-        
-        assert mode in ['image', 'text', 'multimodal'], "mode parameter must be image, text, or multimodal"
-        text = self.tokenizer(caption, return_tensors="pt").to(image.device) 
-        
-        if mode=='image':    
-            # return image features
-            image_embeds = self.visual_encoder(image)             
-            return image_embeds
-        
-        elif mode=='text':
-            # return text features
-            text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      
-                                            return_dict = True, mode = 'text')  
-            return text_output.last_hidden_state
-        
-        elif mode=='multimodal':
-            # return multimodel features
-            image_embeds = self.visual_encoder(image)    
-            image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)      
-            
-            text.input_ids[:,0] = self.tokenizer.enc_token_id
-            output = self.text_encoder(text.input_ids,
-                                       attention_mask = text.attention_mask,
-                                       encoder_hidden_states = image_embeds,
-                                       encoder_attention_mask = image_atts,      
-                                       return_dict = True,
-                                      )              
-            return output.last_hidden_state
-        
-        
-        
-class BLIP_Decoder(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/med_config.json',  
-                 image_size = 384,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,
-                 prompt = 'a picture of ',
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """            
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
-        self.tokenizer = init_tokenizer()   
-        med_config = BertConfig.from_json_file(med_config)
-        med_config.encoder_width = vision_width
-        self.text_decoder = BertLMHeadModel(config=med_config)    
-        
-        self.prompt = prompt
-        self.prompt_length = len(self.tokenizer(self.prompt).input_ids)-1
-
-        
-    def forward(self, image, caption):
-        
-        image_embeds = self.visual_encoder(image) 
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
-        
-        text = self.tokenizer(caption, padding='longest', truncation=True, max_length=40, return_tensors="pt").to(image.device) 
-        
-        text.input_ids[:,0] = self.tokenizer.bos_token_id
-        
-        decoder_targets = text.input_ids.masked_fill(text.input_ids == self.tokenizer.pad_token_id, -100)         
-        decoder_targets[:,:self.prompt_length] = -100
-     
-        decoder_output = self.text_decoder(text.input_ids, 
-                                           attention_mask = text.attention_mask, 
-                                           encoder_hidden_states = image_embeds,
-                                           encoder_attention_mask = image_atts,                  
-                                           labels = decoder_targets,
-                                           return_dict = True,   
-                                          )   
-        loss_lm = decoder_output.loss
-        
-        return loss_lm
-        
-    def generate(self, image, sample=False, num_beams=3, max_length=30, min_length=10, top_p=0.9, repetition_penalty=1.0):
-        image_embeds = self.visual_encoder(image)
-
-        if not sample:
-            image_embeds = image_embeds.repeat_interleave(num_beams,dim=0)
-            
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
-        model_kwargs = {"encoder_hidden_states": image_embeds, "encoder_attention_mask":image_atts}
-        
-        prompt = [self.prompt] * image.size(0)
-        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(image.device) 
-        input_ids[:,0] = self.tokenizer.bos_token_id
-        input_ids = input_ids[:, :-1] 
-
-        if sample:
-            #nucleus sampling
-            outputs = self.text_decoder.generate(input_ids=input_ids,
-                                                  max_length=max_length,
-                                                  min_length=min_length,
-                                                  do_sample=True,
-                                                  top_p=top_p,
-                                                  num_return_sequences=1,
-                                                  eos_token_id=self.tokenizer.sep_token_id,
-                                                  pad_token_id=self.tokenizer.pad_token_id, 
-                                                  repetition_penalty=1.1,                                            
-                                                  **model_kwargs)
-        else:
-            #beam search
-            outputs = self.text_decoder.generate(input_ids=input_ids,
-                                                  max_length=max_length,
-                                                  min_length=min_length,
-                                                  num_beams=num_beams,
-                                                  eos_token_id=self.tokenizer.sep_token_id,
-                                                  pad_token_id=self.tokenizer.pad_token_id,     
-                                                  repetition_penalty=repetition_penalty,
-                                                  **model_kwargs)            
-            
-        captions = []    
-        for output in outputs:
-            caption = self.tokenizer.decode(output, skip_special_tokens=True)    
-            captions.append(caption[len(self.prompt):])
-        return captions
-    
-
-def blip_decoder(pretrained='',**kwargs):
-    model = BLIP_Decoder(**kwargs)
-    if pretrained:
-        model,msg = load_checkpoint(model,pretrained)
-        assert(len(msg.missing_keys)==0)
-    return model    
-    
-def blip_feature_extractor(pretrained='',**kwargs):
-    model = BLIP_Base(**kwargs)
-    if pretrained:
-        model,msg = load_checkpoint(model,pretrained)
-        assert(len(msg.missing_keys)==0)
-    return model        
-
-def init_tokenizer():
-    tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "bert_tokenizer")
-    tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
-    tokenizer.add_special_tokens({'bos_token':'[DEC]'})
-    tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       
-    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  
-    return tokenizer
-
-
-def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
-        
-    assert vit in ['base', 'large'], "vit parameter must be base or large"
-    if vit=='base':
-        vision_width = 768
-        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12, 
-                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
-                                           drop_path_rate=0 or drop_path_rate
-                                          )   
-    elif vit=='large':
-        vision_width = 1024
-        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24, 
-                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
-                                           drop_path_rate=0.1 or drop_path_rate
-                                          )   
-    return visual_encoder, vision_width
-
-def is_url(url_or_filename):
-    parsed = urlparse(url_or_filename)
-    return parsed.scheme in ("http", "https")
-
-def load_checkpoint(model,url_or_filename):
-    if is_url(url_or_filename):
-        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
-        checkpoint = torch.load(cached_file, map_location='cpu') 
-    elif os.path.isfile(url_or_filename):        
-        checkpoint = torch.load(url_or_filename, map_location='cpu') 
-    else:
-        raise RuntimeError('checkpoint url or path is invalid')
-        
-    state_dict = checkpoint['model']
-    
-    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) 
-    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
-        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
-                                                                         model.visual_encoder_m)    
-    for key in model.state_dict().keys():
-        if key in state_dict.keys():
-            if state_dict[key].shape!=model.state_dict()[key].shape:
-                del state_dict[key]
-    
-    msg = model.load_state_dict(state_dict,strict=False)
-    print('load checkpoint from %s'%url_or_filename)  
-    return model,msg
-    

repositories/Fooocus/extras/BLIP/models/blip_itm.py

@@ -1,76 +0,0 @@
-from extras.BLIP.models.med import BertConfig, BertModel
-from transformers import BertTokenizer
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
-
-class BLIP_ITM(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/med_config.json',  
-                 image_size = 384,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,                      
-                 embed_dim = 256,     
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """               
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
-        self.tokenizer = init_tokenizer()   
-        med_config = BertConfig.from_json_file(med_config)
-        med_config.encoder_width = vision_width
-        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)          
-
-        text_width = self.text_encoder.config.hidden_size
-        
-        self.vision_proj = nn.Linear(vision_width, embed_dim)
-        self.text_proj = nn.Linear(text_width, embed_dim)
-
-        self.itm_head = nn.Linear(text_width, 2) 
-        
-        
-    def forward(self, image, caption, match_head='itm'):
-
-        image_embeds = self.visual_encoder(image) 
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)        
-      
-        text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=35, 
-                              return_tensors="pt").to(image.device) 
-
-                 
-        if match_head=='itm':
-            output = self.text_encoder(text.input_ids,
-                                       attention_mask = text.attention_mask,
-                                       encoder_hidden_states = image_embeds,
-                                       encoder_attention_mask = image_atts,      
-                                       return_dict = True,
-                                      )
-            itm_output = self.itm_head(output.last_hidden_state[:,0,:])     
-            return itm_output
-            
-        elif match_head=='itc':
-            text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      
-                                            return_dict = True, mode = 'text')                     
-            image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)   
-            text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)    
-            
-            sim = image_feat @ text_feat.t()
-            return sim
-        
-        
-def blip_itm(pretrained='',**kwargs):
-    model = BLIP_ITM(**kwargs)
-    if pretrained:
-        model,msg = load_checkpoint(model,pretrained)
-        assert(len(msg.missing_keys)==0)
-    return model         
-            

repositories/Fooocus/extras/BLIP/models/blip_nlvr.py

@@ -1,105 +0,0 @@
-from extras.BLIP.models.med import BertConfig
-from extras.BLIP.models.nlvr_encoder import BertModel
-from extras.BLIP.models.vit import interpolate_pos_embed
-from extras.BLIP.models.blip import create_vit, init_tokenizer, is_url
-
-from timm.models.hub import download_cached_file
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-from transformers import BertTokenizer
-import numpy as np
-import os
-
-
-class BLIP_NLVR(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/med_config.json',  
-                 image_size = 480,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,                   
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """               
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, drop_path_rate=0.1)
-        self.tokenizer = init_tokenizer()   
-        med_config = BertConfig.from_json_file(med_config)
-        med_config.encoder_width = vision_width
-        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False) 
-                    
-        self.cls_head = nn.Sequential(
-                  nn.Linear(self.text_encoder.config.hidden_size, self.text_encoder.config.hidden_size),
-                  nn.ReLU(),
-                  nn.Linear(self.text_encoder.config.hidden_size, 2)
-                )  
-
-    def forward(self, image, text, targets, train=True):
-        
-        image_embeds = self.visual_encoder(image) 
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)        
-        image0_embeds, image1_embeds = torch.split(image_embeds,targets.size(0))     
-
-        text = self.tokenizer(text, padding='longest', return_tensors="pt").to(image.device) 
-        text.input_ids[:,0] = self.tokenizer.enc_token_id        
-
-        output = self.text_encoder(text.input_ids, 
-                                   attention_mask = text.attention_mask, 
-                                   encoder_hidden_states = [image0_embeds,image1_embeds],
-                                   encoder_attention_mask = [image_atts[:image0_embeds.size(0)],
-                                                             image_atts[image0_embeds.size(0):]],        
-                                   return_dict = True,
-                                  )  
-        hidden_state = output.last_hidden_state[:,0,:]        
-        prediction = self.cls_head(hidden_state)
-
-        if train:            
-            loss = F.cross_entropy(prediction, targets)   
-            return loss
-        else:
-            return prediction
-    
-def blip_nlvr(pretrained='',**kwargs):
-    model = BLIP_NLVR(**kwargs)
-    if pretrained:
-        model,msg = load_checkpoint(model,pretrained)
-        print("missing keys:")
-        print(msg.missing_keys)
-    return model  
-
-        
-def load_checkpoint(model,url_or_filename):
-    if is_url(url_or_filename):
-        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
-        checkpoint = torch.load(cached_file, map_location='cpu') 
-    elif os.path.isfile(url_or_filename):        
-        checkpoint = torch.load(url_or_filename, map_location='cpu') 
-    else:
-        raise RuntimeError('checkpoint url or path is invalid')
-    state_dict = checkpoint['model']
-    
-    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) 
-    
-    for key in list(state_dict.keys()):
-        if 'crossattention.self.' in key:
-            new_key0 = key.replace('self','self0')
-            new_key1 = key.replace('self','self1')
-            state_dict[new_key0] = state_dict[key]
-            state_dict[new_key1] = state_dict[key]
-        elif 'crossattention.output.dense.' in key:
-            new_key0 = key.replace('dense','dense0')
-            new_key1 = key.replace('dense','dense1')
-            state_dict[new_key0] = state_dict[key]
-            state_dict[new_key1] = state_dict[key]  
-                
-    msg = model.load_state_dict(state_dict,strict=False)
-    print('load checkpoint from %s'%url_or_filename)  
-    return model,msg
-            

repositories/Fooocus/extras/BLIP/models/blip_pretrain.py

@@ -1,339 +0,0 @@
-'''
- * Copyright (c) 2022, salesforce.com, inc.
- * All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
- * By Junnan Li
-'''
-from extras.BLIP.models.med import BertConfig, BertModel, BertLMHeadModel
-from transformers import BertTokenizer
-import transformers
-transformers.logging.set_verbosity_error()
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
-
-class BLIP_Pretrain(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/bert_config.json',  
-                 image_size = 224,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,                    
-                 embed_dim = 256,     
-                 queue_size = 57600,
-                 momentum = 0.995,
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """               
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
-        
-        if vit=='base':
-            checkpoint = torch.hub.load_state_dict_from_url(
-                url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
-                map_location="cpu", check_hash=True)
-            state_dict = checkpoint["model"]     
-            msg = self.visual_encoder.load_state_dict(state_dict,strict=False)
-        elif vit=='large':
-            from timm.models.helpers import load_custom_pretrained
-            from timm.models.vision_transformer import default_cfgs
-            load_custom_pretrained(self.visual_encoder,default_cfgs['vit_large_patch16_224_in21k'])        
-               
-        self.tokenizer = init_tokenizer()   
-        encoder_config = BertConfig.from_json_file(med_config)
-        encoder_config.encoder_width = vision_width
-        self.text_encoder = BertModel.from_pretrained('bert-base-uncased',config=encoder_config, add_pooling_layer=False)
-        self.text_encoder.resize_token_embeddings(len(self.tokenizer)) 
-
-        text_width = self.text_encoder.config.hidden_size
-        
-        self.vision_proj = nn.Linear(vision_width, embed_dim)
-        self.text_proj = nn.Linear(text_width, embed_dim)
-
-        self.itm_head = nn.Linear(text_width, 2) 
-        
-        # create momentum encoders  
-        self.visual_encoder_m, vision_width = create_vit(vit,image_size)              
-        self.vision_proj_m = nn.Linear(vision_width, embed_dim)
-        self.text_encoder_m = BertModel(config=encoder_config, add_pooling_layer=False)      
-        self.text_proj_m = nn.Linear(text_width, embed_dim)
-        
-        self.model_pairs = [[self.visual_encoder,self.visual_encoder_m],
-                            [self.vision_proj,self.vision_proj_m],
-                            [self.text_encoder,self.text_encoder_m],
-                            [self.text_proj,self.text_proj_m],
-                           ]       
-        self.copy_params()
-
-        # create the queue
-        self.register_buffer("image_queue", torch.randn(embed_dim, queue_size))
-        self.register_buffer("text_queue", torch.randn(embed_dim, queue_size))
-        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))  
-
-        self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
-        self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
-        
-        self.queue_size = queue_size
-        self.momentum = momentum
-        self.temp = nn.Parameter(0.07*torch.ones([]))   
-        
-        # create the decoder
-        decoder_config = BertConfig.from_json_file(med_config)
-        decoder_config.encoder_width = vision_width        
-        self.text_decoder = BertLMHeadModel.from_pretrained('bert-base-uncased',config=decoder_config)    
-        self.text_decoder.resize_token_embeddings(len(self.tokenizer)) 
-        tie_encoder_decoder_weights(self.text_encoder,self.text_decoder.bert,'','/attention')
-        
-        
-    def forward(self, image, caption, alpha):
-        with torch.no_grad():
-            self.temp.clamp_(0.001,0.5)
-        
-        image_embeds = self.visual_encoder(image) 
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)        
-        image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)          
-        
-        text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=30, 
-                              return_tensors="pt").to(image.device)  
-        text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      
-                                        return_dict = True, mode = 'text')            
-        text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)                 
-             
-        # get momentum features
-        with torch.no_grad():
-            self._momentum_update()
-            image_embeds_m = self.visual_encoder_m(image) 
-            image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)  
-            image_feat_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)                   
-            
-            text_output_m = self.text_encoder_m(text.input_ids, attention_mask = text.attention_mask,                      
-                                                return_dict = True, mode = 'text')    
-            text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1) 
-            text_feat_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
-
-            sim_i2t_m = image_feat_m @ text_feat_all / self.temp  
-            sim_t2i_m = text_feat_m @ image_feat_all / self.temp 
-
-            sim_targets = torch.zeros(sim_i2t_m.size()).to(image.device)
-            sim_targets.fill_diagonal_(1)          
-
-            sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
-            sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets        
-
-        sim_i2t = image_feat @ text_feat_all / self.temp
-        sim_t2i = text_feat @ image_feat_all / self.temp
-                             
-        loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
-        loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean() 
-
-        loss_ita = (loss_i2t+loss_t2i)/2
-
-        self._dequeue_and_enqueue(image_feat_m, text_feat_m)        
-
-        ###============== Image-text Matching ===================###
-        encoder_input_ids = text.input_ids.clone()
-        encoder_input_ids[:,0] = self.tokenizer.enc_token_id
-        
-        # forward the positve image-text pair
-        bs = image.size(0)
-        output_pos = self.text_encoder(encoder_input_ids,
-                                       attention_mask = text.attention_mask,
-                                       encoder_hidden_states = image_embeds,
-                                       encoder_attention_mask = image_atts,      
-                                       return_dict = True,
-                                      )            
-        with torch.no_grad():       
-            weights_t2i = F.softmax(sim_t2i[:,:bs],dim=1)+1e-4 
-            weights_t2i.fill_diagonal_(0)            
-            weights_i2t = F.softmax(sim_i2t[:,:bs],dim=1)+1e-4  
-            weights_i2t.fill_diagonal_(0)   
-            
-        # select a negative image for each text
-        image_embeds_neg = []    
-        for b in range(bs):
-            neg_idx = torch.multinomial(weights_t2i[b], 1).item()
-            image_embeds_neg.append(image_embeds[neg_idx])
-        image_embeds_neg = torch.stack(image_embeds_neg,dim=0)   
-
-        # select a negative text for each image
-        text_ids_neg = []
-        text_atts_neg = []
-        for b in range(bs):
-            neg_idx = torch.multinomial(weights_i2t[b], 1).item()
-            text_ids_neg.append(encoder_input_ids[neg_idx])
-            text_atts_neg.append(text.attention_mask[neg_idx])
-
-        text_ids_neg = torch.stack(text_ids_neg,dim=0)   
-        text_atts_neg = torch.stack(text_atts_neg,dim=0)      
-
-        text_ids_all = torch.cat([encoder_input_ids, text_ids_neg],dim=0)     
-        text_atts_all = torch.cat([text.attention_mask, text_atts_neg],dim=0)     
-
-        image_embeds_all = torch.cat([image_embeds_neg,image_embeds],dim=0)
-        image_atts_all = torch.cat([image_atts,image_atts],dim=0)
-
-        output_neg = self.text_encoder(text_ids_all,
-                                       attention_mask = text_atts_all,
-                                       encoder_hidden_states = image_embeds_all,
-                                       encoder_attention_mask = image_atts_all,      
-                                       return_dict = True,
-                                      )                            
-
-        vl_embeddings = torch.cat([output_pos.last_hidden_state[:,0,:], output_neg.last_hidden_state[:,0,:]],dim=0)
-        vl_output = self.itm_head(vl_embeddings)            
-
-        itm_labels = torch.cat([torch.ones(bs,dtype=torch.long),torch.zeros(2*bs,dtype=torch.long)],
-                               dim=0).to(image.device)
-        loss_itm = F.cross_entropy(vl_output, itm_labels)  
-        
-        ##================= LM ========================##     
-        decoder_input_ids = text.input_ids.clone()      
-        decoder_input_ids[:,0] = self.tokenizer.bos_token_id
-        decoder_targets = decoder_input_ids.masked_fill(decoder_input_ids == self.tokenizer.pad_token_id, -100) 
-
-        decoder_output = self.text_decoder(decoder_input_ids, 
-                                           attention_mask = text.attention_mask, 
-                                           encoder_hidden_states = image_embeds,
-                                           encoder_attention_mask = image_atts,                  
-                                           labels = decoder_targets,
-                                           return_dict = True,   
-                                          )   
-          
-        loss_lm = decoder_output.loss                
-        return loss_ita, loss_itm, loss_lm
- 
-
-
-    @torch.no_grad()    
-    def copy_params(self):
-        for model_pair in self.model_pairs:           
-            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
-                param_m.data.copy_(param.data)  # initialize
-                param_m.requires_grad = False  # not update by gradient    
-
-            
-    @torch.no_grad()        
-    def _momentum_update(self):
-        for model_pair in self.model_pairs:           
-            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
-                param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
-
-                        
-    @torch.no_grad()
-    def _dequeue_and_enqueue(self, image_feat, text_feat):
-        # gather keys before updating queue
-        image_feats = concat_all_gather(image_feat)
-        text_feats = concat_all_gather(text_feat)
-
-        batch_size = image_feats.shape[0]
-
-        ptr = int(self.queue_ptr)
-        assert self.queue_size % batch_size == 0  # for simplicity
-
-        # replace the keys at ptr (dequeue and enqueue)
-        self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
-        self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
-        ptr = (ptr + batch_size) % self.queue_size  # move pointer
-
-        self.queue_ptr[0] = ptr 
-
-
-def blip_pretrain(**kwargs):
-    model = BLIP_Pretrain(**kwargs)
-    return model 
-
-
-@torch.no_grad()
-def concat_all_gather(tensor):
-    """
-    Performs all_gather operation on the provided tensors.
-    *** Warning ***: torch.distributed.all_gather has no gradient.
-    """
-    tensors_gather = [torch.ones_like(tensor)
-        for _ in range(torch.distributed.get_world_size())]
-    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
-
-    output = torch.cat(tensors_gather, dim=0)
-    return output     
-
-
-from typing import List
-def tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str, skip_key:str):
-    uninitialized_encoder_weights: List[str] = []
-    if decoder.__class__ != encoder.__class__:
-        print(
-            f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder weights are correctly initialized."
-        )
-
-    def tie_encoder_to_decoder_recursively(
-        decoder_pointer: nn.Module,
-        encoder_pointer: nn.Module,
-        module_name: str,
-        uninitialized_encoder_weights: List[str],
-        skip_key: str,
-        depth=0,
-    ):
-        assert isinstance(decoder_pointer, nn.Module) and isinstance(
-            encoder_pointer, nn.Module
-        ), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module"
-        if hasattr(decoder_pointer, "weight") and skip_key not in module_name:
-            assert hasattr(encoder_pointer, "weight")
-            encoder_pointer.weight = decoder_pointer.weight
-            if hasattr(decoder_pointer, "bias"):
-                assert hasattr(encoder_pointer, "bias")
-                encoder_pointer.bias = decoder_pointer.bias                
-            print(module_name+' is tied')    
-            return
-
-        encoder_modules = encoder_pointer._modules
-        decoder_modules = decoder_pointer._modules
-        if len(decoder_modules) > 0:
-            assert (
-                len(encoder_modules) > 0
-            ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"
-
-            all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
-            encoder_layer_pos = 0
-            for name, module in decoder_modules.items():
-                if name.isdigit():
-                    encoder_name = str(int(name) + encoder_layer_pos)
-                    decoder_name = name
-                    if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
-                        encoder_modules
-                    ) != len(decoder_modules):
-                        # this can happen if the name corresponds to the position in a list module list of layers
-                        # in this case the decoder has added a cross-attention that the encoder does not have
-                        # thus skip this step and subtract one layer pos from encoder
-                        encoder_layer_pos -= 1
-                        continue
-                elif name not in encoder_modules:
-                    continue
-                elif depth > 500:
-                    raise ValueError(
-                        "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model."
-                    )
-                else:
-                    decoder_name = encoder_name = name
-                tie_encoder_to_decoder_recursively(
-                    decoder_modules[decoder_name],
-                    encoder_modules[encoder_name],
-                    module_name + "/" + name,
-                    uninitialized_encoder_weights,
-                    skip_key,
-                    depth=depth + 1,
-                )
-                all_encoder_weights.remove(module_name + "/" + encoder_name)
-
-            uninitialized_encoder_weights += list(all_encoder_weights)
-
-    # tie weights recursively
-    tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights, skip_key)  

repositories/Fooocus/extras/BLIP/models/blip_retrieval.py

@@ -1,319 +0,0 @@
-from extras.BLIP.models.med import BertConfig, BertModel
-from transformers import BertTokenizer
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
-
-class BLIP_Retrieval(nn.Module):
-    def __init__(self,                 
-                 med_config = 'configs/med_config.json',  
-                 image_size = 384,
-                 vit = 'base',
-                 vit_grad_ckpt = False,
-                 vit_ckpt_layer = 0,                      
-                 embed_dim = 256,     
-                 queue_size = 57600,
-                 momentum = 0.995,
-                 negative_all_rank = False,
-                 ):
-        """
-        Args:
-            med_config (str): path for the mixture of encoder-decoder model's configuration file
-            image_size (int): input image size
-            vit (str): model size of vision transformer
-        """               
-        super().__init__()
-        
-        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
-        self.tokenizer = init_tokenizer()   
-        med_config = BertConfig.from_json_file(med_config)
-        med_config.encoder_width = vision_width
-        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)          
-
-        text_width = self.text_encoder.config.hidden_size
-        
-        self.vision_proj = nn.Linear(vision_width, embed_dim)
-        self.text_proj = nn.Linear(text_width, embed_dim)
-
-        self.itm_head = nn.Linear(text_width, 2) 
-        
-        # create momentum encoders  
-        self.visual_encoder_m, vision_width = create_vit(vit,image_size)              
-        self.vision_proj_m = nn.Linear(vision_width, embed_dim)
-        self.text_encoder_m = BertModel(config=med_config, add_pooling_layer=False)    
-        self.text_proj_m = nn.Linear(text_width, embed_dim)
-        
-        self.model_pairs = [[self.visual_encoder,self.visual_encoder_m],
-                            [self.vision_proj,self.vision_proj_m],
-                            [self.text_encoder,self.text_encoder_m],
-                            [self.text_proj,self.text_proj_m],
-                           ]       
-        self.copy_params()
-
-        # create the queue
-        self.register_buffer("image_queue", torch.randn(embed_dim, queue_size))
-        self.register_buffer("text_queue", torch.randn(embed_dim, queue_size))
-        self.register_buffer("idx_queue", torch.full((1,queue_size),-100))
-        self.register_buffer("ptr_queue", torch.zeros(1, dtype=torch.long))  
-
-        self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
-        self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
-        
-        self.queue_size = queue_size
-        self.momentum = momentum
-        self.temp = nn.Parameter(0.07*torch.ones([]))   
-        
-        self.negative_all_rank = negative_all_rank
-        
-        
-    def forward(self, image, caption, alpha, idx):
-        with torch.no_grad():
-            self.temp.clamp_(0.001,0.5)
-        
-        image_embeds = self.visual_encoder(image) 
-        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)        
-        image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)    
-        
-        text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=35, 
-                              return_tensors="pt").to(image.device) 
-        
-        text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      
-                                        return_dict = True, mode = 'text')            
-        text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)        
-        
-        ###============== Image-text Contrastive Learning ===================###
-        idx = idx.view(-1,1)
-        idx_all = torch.cat([idx.t(), self.idx_queue.clone().detach()],dim=1)  
-        pos_idx = torch.eq(idx, idx_all).float()       
-        sim_targets = pos_idx / pos_idx.sum(1,keepdim=True)   
-        
-        # get momentum features
-        with torch.no_grad():
-            self._momentum_update()
-            image_embeds_m = self.visual_encoder_m(image) 
-            image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)  
-            image_feat_m_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)                   
-            
-            text_output_m = self.text_encoder_m(text.input_ids, attention_mask = text.attention_mask,                      
-                                                return_dict = True, mode = 'text')    
-            text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1) 
-            text_feat_m_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
-
-            sim_i2t_m = image_feat_m @ text_feat_m_all / self.temp  
-            sim_t2i_m = text_feat_m @ image_feat_m_all / self.temp   
-
-            sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
-            sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets        
-
-        sim_i2t = image_feat @ text_feat_m_all / self.temp 
-        sim_t2i = text_feat @ image_feat_m_all / self.temp 
-                             
-        loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
-        loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean() 
-
-        loss_ita = (loss_i2t+loss_t2i)/2
-        
-        idxs = concat_all_gather(idx)
-        self._dequeue_and_enqueue(image_feat_m, text_feat_m, idxs)        
-
-        ###============== Image-text Matching ===================###
-        encoder_input_ids = text.input_ids.clone()
-        encoder_input_ids[:,0] = self.tokenizer.enc_token_id
-
-        # forward the positve image-text pair
-        bs = image.size(0)
-        output_pos = self.text_encoder(encoder_input_ids,
-                                       attention_mask = text.attention_mask,
-                                       encoder_hidden_states = image_embeds,
-                                       encoder_attention_mask = image_atts,      
-                                       return_dict = True,
-                                      )  
-        
-        
-        if self.negative_all_rank:    
-            # compute sample similarity
-            with torch.no_grad():                
-                mask = torch.eq(idx, idxs.t())
-
-                image_feat_world = concat_all_gather(image_feat)
-                text_feat_world = concat_all_gather(text_feat)
-
-                sim_i2t = image_feat @ text_feat_world.t() / self.temp 
-                sim_t2i = text_feat @ image_feat_world.t() / self.temp 
-
-                weights_i2t = F.softmax(sim_i2t,dim=1)
-                weights_i2t.masked_fill_(mask, 0)            
-
-                weights_t2i = F.softmax(sim_t2i,dim=1)
-                weights_t2i.masked_fill_(mask, 0)     
-
-            image_embeds_world = all_gather_with_grad(image_embeds) 
-
-            # select a negative image (from all ranks) for each text
-            image_embeds_neg = []    
-            for b in range(bs):
-                neg_idx = torch.multinomial(weights_t2i[b], 1).item()
-                image_embeds_neg.append(image_embeds_world[neg_idx])
-            image_embeds_neg = torch.stack(image_embeds_neg,dim=0)   
-
-            # select a negative text (from all ranks) for each image
-            input_ids_world = concat_all_gather(encoder_input_ids)
-            att_mask_world = concat_all_gather(text.attention_mask)        
-
-            text_ids_neg = []
-            text_atts_neg = []
-            for b in range(bs):
-                neg_idx = torch.multinomial(weights_i2t[b], 1).item()
-                text_ids_neg.append(input_ids_world[neg_idx])
-                text_atts_neg.append(att_mask_world[neg_idx])
-                
-        else:
-            with torch.no_grad():                
-                mask = torch.eq(idx, idx.t())
-                
-                sim_i2t = image_feat @ text_feat.t() / self.temp 
-                sim_t2i = text_feat @ image_feat.t() / self.temp 
-
-                weights_i2t = F.softmax(sim_i2t,dim=1)
-                weights_i2t.masked_fill_(mask, 0)            
-
-                weights_t2i = F.softmax(sim_t2i,dim=1)
-                weights_t2i.masked_fill_(mask, 0)     
-
-            # select a negative image (from same rank) for each text
-            image_embeds_neg = []    
-            for b in range(bs):
-                neg_idx = torch.multinomial(weights_t2i[b], 1).item()
-                image_embeds_neg.append(image_embeds[neg_idx])
-            image_embeds_neg = torch.stack(image_embeds_neg,dim=0)   
-
-            # select a negative text (from same rank) for each image    
-            text_ids_neg = []
-            text_atts_neg = []
-            for b in range(bs):
-                neg_idx = torch.multinomial(weights_i2t[b], 1).item()
-                text_ids_neg.append(encoder_input_ids[neg_idx])
-                text_atts_neg.append(text.attention_mask[neg_idx])            
-            
-        text_ids_neg = torch.stack(text_ids_neg,dim=0)   
-        text_atts_neg = torch.stack(text_atts_neg,dim=0)      
-
-        text_ids_all = torch.cat([encoder_input_ids, text_ids_neg],dim=0)     
-        text_atts_all = torch.cat([text.attention_mask, text_atts_neg],dim=0)     
-
-        image_embeds_all = torch.cat([image_embeds_neg,image_embeds],dim=0)
-        image_atts_all = torch.cat([image_atts,image_atts],dim=0)
-
-        output_neg = self.text_encoder(text_ids_all,
-                                       attention_mask = text_atts_all,
-                                       encoder_hidden_states = image_embeds_all,
-                                       encoder_attention_mask = image_atts_all,      
-                                       return_dict = True,
-                                      )                         
-          
-
-        vl_embeddings = torch.cat([output_pos.last_hidden_state[:,0,:], output_neg.last_hidden_state[:,0,:]],dim=0)
-        vl_output = self.itm_head(vl_embeddings)            
-
-        itm_labels = torch.cat([torch.ones(bs,dtype=torch.long),torch.zeros(2*bs,dtype=torch.long)],
-                               dim=0).to(image.device)
-        loss_itm = F.cross_entropy(vl_output, itm_labels)     
-
-        return loss_ita, loss_itm 
- 
-
-    @torch.no_grad()    
-    def copy_params(self):
-        for model_pair in self.model_pairs:           
-            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
-                param_m.data.copy_(param.data)  # initialize
-                param_m.requires_grad = False  # not update by gradient    
-
-            
-    @torch.no_grad()        
-    def _momentum_update(self):
-        for model_pair in self.model_pairs:           
-            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
-                param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
-                
-                
-    @torch.no_grad()
-    def _dequeue_and_enqueue(self, image_feat, text_feat, idxs):
-        # gather keys before updating queue
-        image_feats = concat_all_gather(image_feat)
-        text_feats = concat_all_gather(text_feat)
-        
-
-        batch_size = image_feats.shape[0]
-
-        ptr = int(self.ptr_queue)
-        assert self.queue_size % batch_size == 0  # for simplicity
-
-        # replace the keys at ptr (dequeue and enqueue)
-        self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
-        self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
-        self.idx_queue[:, ptr:ptr + batch_size] = idxs.T
-        ptr = (ptr + batch_size) % self.queue_size # move pointer
-
-        self.ptr_queue[0] = ptr  
-
-
-def blip_retrieval(pretrained='',**kwargs):
-    model = BLIP_Retrieval(**kwargs)
-    if pretrained:
-        model,msg = load_checkpoint(model,pretrained)
-        print("missing keys:")
-        print(msg.missing_keys)
-    return model 
-
-
-@torch.no_grad()
-def concat_all_gather(tensor):
-    """
-    Performs all_gather operation on the provided tensors.
-    *** Warning ***: torch.distributed.all_gather has no gradient.
-    """
-    tensors_gather = [torch.ones_like(tensor)
-        for _ in range(torch.distributed.get_world_size())]
-    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
-
-    output = torch.cat(tensors_gather, dim=0)
-    return output      
-
-
-class GatherLayer(torch.autograd.Function):
-    """
-    Gather tensors from all workers with support for backward propagation:
-    This implementation does not cut the gradients as torch.distributed.all_gather does.
-    """
-
-    @staticmethod
-    def forward(ctx, x):
-        output = [torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())]
-        torch.distributed.all_gather(output, x)
-        return tuple(output)
-
-    @staticmethod
-    def backward(ctx, *grads):
-        all_gradients = torch.stack(grads)
-        torch.distributed.all_reduce(all_gradients)
-        return all_gradients[torch.distributed.get_rank()]
-
-
-def all_gather_with_grad(tensors):
-    """
-    Performs all_gather operation on the provided tensors.
-    Graph remains connected for backward grad computation.
-    """
-    # Queue the gathered tensors
-    world_size = torch.distributed.get_world_size()
-    # There is no need for reduction in the single-proc case
-    if world_size == 1:
-        return tensors
-
-    tensor_all = GatherLayer.apply(tensors)
-
-    return torch.cat(tensor_all, dim=0)