Merge pull request #9 from konieshadow/cog

Add cog config for replicate.com

.gitignore

@@ -38,3 +38,4 @@ notification.mp3
 /node_modules
 /package-lock.json
 /.coverage*
+.cog/

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.
 

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 <package-name>==<version>
+  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"

fooocus_api_version.py

@@ -1 +1 @@
-version = '0.1.10'
+version = '0.1.11'

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)
 
 

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')

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

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

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)
+            self.history.append(task)