diff --git a/.gitignore b/.gitignore index 880f47f..0a95e53 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,4 @@ notification.mp3 /node_modules /package-lock.json /.coverage* +.cog/ \ No newline at end of file diff --git a/README.md b/README.md index f3c41a0..d95f40d 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ Currently loaded Fooocus version: 2.1.25 Need python version >= 3.10 ``` pip install -r requirements.txt -pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118 xformers +pip install torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/cu118 xformers ``` You may change the part "cu118" of extra-index-url to your local installed cuda driver version. diff --git a/cog.yaml b/cog.yaml new file mode 100644 index 0000000..89b8512 --- /dev/null +++ b/cog.yaml @@ -0,0 +1,46 @@ +# Configuration for Cog ⚙️ +# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md + +build: + # set to true if your model requires a GPU + gpu: true + cuda: "11.8" + + # a list of ubuntu apt packages to install + system_packages: + - "libgl1-mesa-glx" + - "libglib2.0-0" + + # python version in the form '3.11' or '3.11.4' + python_version: "3.11" + + # a list of packages in the format == + python_packages: + - "torchsde==0.2.5" + - "einops==0.4.1" + - "transformers==4.30.2" + - "safetensors==0.3.1" + - "accelerate==0.21.0" + - "pyyaml==6.0" + - "Pillow==9.2.0" + - "scipy==1.9.3" + - "tqdm==4.64.1" + - "psutil==5.9.5" + - "numpy==1.23.5" + - "pytorch_lightning==1.9.4" + - "omegaconf==2.2.3" + - "pygit2==1.12.2" + - "opencv-contrib-python==4.8.0.74" + - "torch==2.0.1" + - "torchvision==0.15.2" + - "xformers==0.0.21" + + # commands run after the environment is setup + # run: + # - "echo env is ready!" + # - "echo another command if needed" + +image: "r8.im/konieshadow/fooocus-api" + +# predict.py defines how predictions are run on your model +predict: "predict.py:Predictor" diff --git a/fooocus_api_version.py b/fooocus_api_version.py index 85b15ca..83a3124 100644 --- a/fooocus_api_version.py +++ b/fooocus_api_version.py @@ -1 +1 @@ -version = '0.1.10' \ No newline at end of file +version = '0.1.11' \ No newline at end of file diff --git a/fooocusapi/api.py b/fooocusapi/api.py index ac7bfde..ae23c00 100644 --- a/fooocusapi/api.py +++ b/fooocusapi/api.py @@ -2,7 +2,7 @@ from fastapi import Depends, FastAPI, Header, Query, UploadFile from fastapi.params import File import uvicorn -from fooocusapi.api_utils import generation_output +from fooocusapi.api_utils import generation_output, req_to_params from fooocusapi.models import GeneratedImageBase64, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest from fooocusapi.task_queue import TaskQueue from fooocusapi.worker import process_generate @@ -43,7 +43,7 @@ def text2img_generation(req: Text2ImgRequest, accept: str = Header(None), else: streaming_output = False - results = process_generate(req) + results = process_generate(req_to_params(req)) return generation_output(results, streaming_output) @@ -61,7 +61,7 @@ def img_upscale_or_vary(input_image: UploadFile, req: ImgUpscaleOrVaryRequest = else: streaming_output = False - results = process_generate(req) + results = process_generate(req_to_params(req)) return generation_output(results, streaming_output) @@ -79,7 +79,7 @@ def img_inpaint_or_outpaint(input_image: UploadFile, req: ImgInpaintOrOutpaintRe else: streaming_output = False - results = process_generate(req) + results = process_generate(req_to_params(req)) return generation_output(results, streaming_output) @@ -98,7 +98,7 @@ def img_prompt(cn_img1: Optional[UploadFile] = File(None), else: streaming_output = False - results = process_generate(req) + results = process_generate(req_to_params(req)) return generation_output(results, streaming_output) diff --git a/fooocusapi/api_utils.py b/fooocusapi/api_utils.py index d817e04..751862b 100644 --- a/fooocusapi/api_utils.py +++ b/fooocusapi/api_utils.py @@ -1,15 +1,14 @@ import base64 -import inspect import io from io import BytesIO -from typing import Annotated, List +from typing import List import numpy as np -from fastapi import Form, Response, UploadFile +from fastapi import Response, UploadFile from PIL import Image -from fooocusapi.models import GeneratedImage, GeneratedImageBase64, GenerationFinishReason - -from modules.util import HWC3 +from fooocusapi.models import GeneratedImageBase64, GenerationFinishReason, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest +from fooocusapi.parameters import ImageGenerationParams, ImageGenerationResult +import modules.flags as flags def narray_to_base64img(narray: np.ndarray) -> str: @@ -42,9 +41,69 @@ def read_input_image(input_image: UploadFile) -> np.ndarray: return image -def generation_output(results: List[GeneratedImage], streaming_output: bool) -> Response | List[GeneratedImageBase64]: +def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams: + prompt = req.prompt + negative_prompt = req.negative_promit + style_selections = [s.value for s in req.style_selections] + performance_selection = req.performance_selection.value + aspect_ratios_selection = req.aspect_ratios_selection.value + 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 + loras = [(lora.model_name, lora.weight) for lora in req.loras] + uov_input_image = None if not isinstance( + req, ImgUpscaleOrVaryRequest) else read_input_image(req.input_image) + uov_method = flags.disabled if not isinstance( + req, ImgUpscaleOrVaryRequest) else req.uov_method.value + outpaint_selections = [] if not isinstance(req, ImgInpaintOrOutpaintRequest) else [ + s.value for s in req.outpaint_selections] + + inpaint_input_image = None + if isinstance(req, ImgInpaintOrOutpaintRequest): + input_image = read_input_image(req.input_image) + if req.input_mask is not None: + input_mask = read_input_image(req.input_mask) + else: + input_mask = np.zeros(input_image.shape) + inpaint_input_image = { + 'image': input_image, + 'mask': input_mask + } + + image_prompts = [] + if isinstance(req, ImgPromptRequest): + for img_prompt in req.image_prompts: + if img_prompt.cn_img is not None: + cn_img = read_input_image(img_prompt.cn_img) + image_prompts.append( + (cn_img, img_prompt.cn_stop, img_prompt.cn_weight, img_prompt.cn_type.value)) + + return ImageGenerationParams(prompt=prompt, + negative_promit=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, + loras=loras, + uov_input_image=uov_input_image, + uov_method=uov_method, + outpaint_selections=outpaint_selections, + inpaint_input_image=inpaint_input_image, + image_prompts=image_prompts + ) + + +def generation_output(results: List[ImageGenerationResult], streaming_output: bool) -> Response | List[GeneratedImageBase64]: if streaming_output: - if len(results) == 0 or results[0].finish_reason is not GenerationFinishReason.success: + if len(results) == 0 or results[0].finish_reason != GenerationFinishReason.success: return Response(status_code=500) bytes = narray_to_bytesimg(results[0].im) return Response(bytes, media_type='image/png') diff --git a/fooocusapi/models.py b/fooocusapi/models.py index 8db63a7..8ed8f0e 100644 --- a/fooocusapi/models.py +++ b/fooocusapi/models.py @@ -6,6 +6,7 @@ from enum import Enum from pydantic_core import InitErrorDetails +from fooocusapi.parameters import GenerationFinishReason import modules.flags as flags @@ -539,39 +540,8 @@ def as_form(cls, cn_img1: UploadFile = Form(File(None), description="Input image loras=loras) -class GenerationFinishReason(str, Enum): - success = 'SUCCESS' - queue_is_full = 'QUEUE_IS_FULL' - user_cancel = 'USER_CANCEL' - error = 'ERROR' - - -class GeneratedImage(BaseModel): - im: object | None - seed: int - finish_reason: GenerationFinishReason - - class GeneratedImageBase64(BaseModel): base64: str | None = Field( description="Image encoded in base64, or null if finishReasen is not 'SUCCESS'") seed: int = Field(description="The seed associated with this image") finish_reason: GenerationFinishReason - - -class TaskType(str, Enum): - text2img = 'text2img' - - -class QueueTask(object): - is_finished: bool = False - start_millis: int = 0 - finish_millis: int = 0 - finish_with_error: bool = False - task_result: any = None - - def __init__(self, seq: int, type: TaskType, req_param: dict, in_queue_millis: int): - self.seq = seq - self.type = type - self.req_param = req_param - self.in_queue_millis = in_queue_millis diff --git a/fooocusapi/parameters.py b/fooocusapi/parameters.py new file mode 100644 index 0000000..a756530 --- /dev/null +++ b/fooocusapi/parameters.py @@ -0,0 +1,54 @@ +from enum import Enum +from typing import BinaryIO, Dict, List, Tuple +import numpy as np + + +class GenerationFinishReason(str, Enum): + success = 'SUCCESS' + queue_is_full = 'QUEUE_IS_FULL' + user_cancel = 'USER_CANCEL' + error = 'ERROR' + + +class ImageGenerationResult(object): + def __init__(self, im: np.ndarray | None, seed: int, finish_reason: GenerationFinishReason): + self.im = im + self.seed = seed + self.finish_reason = finish_reason + + +class ImageGenerationParams(object): + def __init__(self, prompt: str, + negative_promit: str, + style_selections: List[str], + performance_selection: List[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, + loras: List[Tuple[str, float]], + uov_input_image: BinaryIO | None, + uov_method: str, + outpaint_selections: List[str], + inpaint_input_image: Dict[str, np.ndarray] | None, + image_prompts: List[Tuple[BinaryIO, float, float, str]]): + self.prompt = prompt + self.negative_promit = negative_promit + 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.loras = loras + self.uov_input_image = uov_input_image + self.uov_method = uov_method + self.outpaint_selections = outpaint_selections + self.inpaint_input_image = inpaint_input_image + self.image_prompts = image_prompts diff --git a/fooocusapi/task_queue.py b/fooocusapi/task_queue.py index 0e1198a..86ade16 100644 --- a/fooocusapi/task_queue.py +++ b/fooocusapi/task_queue.py @@ -1,6 +1,24 @@ +from enum import Enum import time from typing import List -from fooocusapi.models import QueueTask, TaskType + + +class TaskType(str, Enum): + text2img = 'text2img' + + +class QueueTask(object): + is_finished: bool = False + start_millis: int = 0 + finish_millis: int = 0 + finish_with_error: bool = False + task_result: any = None + + def __init__(self, seq: int, type: TaskType, req_param: dict, in_queue_millis: int): + self.seq = seq + self.type = type + self.req_param = req_param + self.in_queue_millis = in_queue_millis class TaskQueue(object): @@ -30,20 +48,19 @@ def get_task(self, seq: int, include_history: bool = False) -> QueueTask | None: if task.seq == seq: return task - if include_history: + if include_history: for task in self.history: if task.seq == seq: return task return None - + def is_task_ready_to_start(self, seq: int) -> bool: task = self.get_task(seq) if task is None: return False - - return self.queue[0].seq == seq + return self.queue[0].seq == seq def start_task(self, seq: int): task = self.get_task(seq) @@ -60,4 +77,4 @@ def finish_task(self, seq: int, task_result: any, finish_with_error: bool): # Move task to history self.queue.remove(task) - self.history.append(task) \ No newline at end of file + self.history.append(task) diff --git a/fooocusapi/worker.py b/fooocusapi/worker.py index e3ae857..bc71412 100644 --- a/fooocusapi/worker.py +++ b/fooocusapi/worker.py @@ -4,16 +4,15 @@ import numpy as np import torch from typing import List -from fooocusapi.api_utils import read_input_image -from fooocusapi.models import GeneratedImage, GenerationFinishReason, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, PerfomanceSelection, TaskType, Text2ImgRequest -from fooocusapi.task_queue import TaskQueue +from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams, ImageGenerationResult +from fooocusapi.task_queue import TaskQueue, TaskType task_queue = TaskQueue() @torch.no_grad() @torch.inference_mode() -def process_generate(req: Text2ImgRequest) -> List[GeneratedImage]: +def process_generate(params: ImageGenerationParams) -> List[ImageGenerationResult]: import modules.default_pipeline as pipeline import modules.patch as patch import modules.flags as flags @@ -37,19 +36,19 @@ def progressbar(number, text): outputs.append(['preview', (number, text, None)]) def make_results_from_outputs(): - results: List[GeneratedImage] = [] + results: List[ImageGenerationResult] = [] for item in outputs: if item[0] == 'results': for im in item[1]: if isinstance(im, np.ndarray): - results.append(GeneratedImage(im=im, seed=item[2], finish_reason=GenerationFinishReason.success)) + results.append(ImageGenerationResult(im=im, seed=item[2], finish_reason=GenerationFinishReason.success)) return results task_seq = task_queue.add_task(TaskType.text2img, { - 'body': req.__dict__}) + 'body': params.__dict__}) if task_seq is None: print("[Task Queue] The task queue has reached limit") - results = [GeneratedImage(im=None, seed=0, + results = [ImageGenerationResult(im=None, seed=0, finish_reason=GenerationFinishReason.queue_is_full)] return results @@ -75,41 +74,29 @@ def make_results_from_outputs(): execution_start_time = time.perf_counter() # Transform pamameters - prompt = req.prompt - negative_prompt = req.negative_promit - style_selections = [s.value for s in req.style_selections] - performance_selection = req.performance_selection.value - aspect_ratios_selection = req.aspect_ratios_selection.value - 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 - loras = [(lora.model_name, lora.weight) for lora in req.loras] - input_image_checkbox = isinstance(req, ImgUpscaleOrVaryRequest) or isinstance(req, ImgInpaintOrOutpaintRequest) or isinstance(req, ImgPromptRequest) - current_tab = 'uov' if isinstance(req, ImgUpscaleOrVaryRequest) else 'inpaint' if isinstance(req, ImgInpaintOrOutpaintRequest) else 'ip' if isinstance(req, ImgPromptRequest) else None - uov_method = flags.disabled if not isinstance(req, ImgUpscaleOrVaryRequest) else req.uov_method.value - uov_input_image = None if not isinstance(req, ImgUpscaleOrVaryRequest) else read_input_image(req.input_image) - outpaint_selections = [] if not isinstance(req, ImgInpaintOrOutpaintRequest) else [s.value for s in req.outpaint_selections] - - inpaint_input_image = None - if isinstance(req, ImgInpaintOrOutpaintRequest): - input_image = read_input_image(req.input_image) - if req.input_mask is not None: - input_mask = read_input_image(req.input_mask) - else: - input_mask = np.zeros(input_image.shape) - inpaint_input_image = { - 'image': input_image, - 'mask': input_mask - } + prompt = params.prompt + negative_prompt = params.negative_promit + style_selections = params.style_selections + performance_selection = params.performance_selection + aspect_ratios_selection = params.aspect_ratios_selection + image_number = params.image_number + image_seed = params.image_seed + sharpness = params.sharpness + guidance_scale = params.guidance_scale + base_model_name = params.base_model_name + refiner_model_name = params.refiner_model_name + 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 'inpaint' if params.inpaint_input_image is not None else 'ip' if len(params.image_prompts) > 0 else None + uov_method = params.uov_method + uov_input_image = params.uov_input_image + outpaint_selections = params.outpaint_selections + inpaint_input_image = params.inpaint_input_image cn_tasks = {flags.cn_ip: [], flags.cn_canny: [], flags.cn_cpds: []} - if isinstance(req, ImgPromptRequest): - for img_prompt in req.image_prompts: - if img_prompt.cn_img is not None: - cn_tasks[img_prompt.cn_type.value].append([read_input_image(img_prompt.cn_img), img_prompt.cn_stop, img_prompt.cn_weight]) + 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]) def build_advanced_parameters(): adm_scaler_positive=1.5 @@ -598,16 +585,16 @@ def callback(step, x0, x, total_steps, y): # Fooocus async_worker.py code end - results.append(GeneratedImage( + results.append(ImageGenerationResult( im=imgs[0], seed=task['task_seed'], finish_reason=GenerationFinishReason.success)) except model_management.InterruptProcessingException as e: print('User stopped') - results.append(GeneratedImage( + results.append(ImageGenerationResult( im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.user_cancel)) break except Exception as e: print('Process failed:', e) - results.append(GeneratedImage( + results.append(ImageGenerationResult( im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.error)) execution_time = time.perf_counter() - execution_start_time diff --git a/predict.py b/predict.py new file mode 100644 index 0000000..4d1a655 --- /dev/null +++ b/predict.py @@ -0,0 +1,87 @@ +# Prediction interface for Cog ⚙️ +# https://github.com/replicate/cog/blob/main/docs/python.md + +import os +from typing import List +from cog import BasePredictor, Input, Path + +from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams +from fooocusapi.worker import process_generate +from PIL import Image + + +class Args(object): + sync_repo = None + + +class Predictor(BasePredictor): + def setup(self) -> None: + """Load the model into memory to make running multiple predictions efficient""" + from main import prepare_environments + print("[Predictor Setup] Prepare environments") + prepare_environments(Args()) + + print("[Predictor Setup] Preload pipeline") + import modules.default_pipeline as _ + print("[Predictor Setup] Finished") + + def predict( + self, + prompt: str = Input( + default='', description="Prompt for image generation") + ) -> List[Path]: + """Run a single prediction on the model""" + from modules.util import generate_temp_filename + import modules.flags as flags + + negative_promit = '' + style_selections = ['Fooocus V2', 'Default (Slightly Cinematic)'] + performance_selection = 'Spped' + aspect_ratios_selection = '1152×896' + image_number = 1 + image_seed = -1 + sharpness = 2.0 + guidance_scale = 7.0 + base_model_name = 'sd_xl_base_1.0_0.9vae.safetensors' + refiner_model_name = 'sd_xl_refiner_1.0_0.9vae.safetensors' + loras = [('sd_xl_offset_example-lora_1.0.safetensors', 0.5)] + uov_input_image = None + uov_method = flags.disabled + outpaint_selections = [] + inpaint_input_image = None + image_prompts = [] + + params = ImageGenerationParams(prompt=prompt, + negative_promit=negative_promit, + 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, + loras=loras, + uov_input_image=uov_input_image, + uov_method=uov_method, + outpaint_selections=outpaint_selections, + inpaint_input_image=inpaint_input_image, + image_prompts=image_prompts + ) + + print(f"[Predictor Predict] Params: {params.__dict__}") + + results = process_generate(params) + + output_paths: List[Path] = [] + for r in results: + if r.finish_reason == GenerationFinishReason.success and r.im is not None: + _, local_temp_filename, _ = generate_temp_filename('/tmp') + os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True) + Image.fromarray(r.im).save(local_temp_filename) + output_paths.append(Path(local_temp_filename)) + + print(f"[Predictor Predict] Finished with {len(output_paths)} images") + + return output_paths