diff --git a/README.md b/README.md index e57d2b8..81aab68 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,6 @@ conda env create -f environment.yaml conda activate fooocus-api ``` -Set enviroment variable `TORCH_INDEX_URL` to the version corresponding to the local cuda driver. -Default is "https://download.pytorch.org/whl/cu121", you may change the part "cu118". - Run ``` python main.py @@ -70,13 +67,15 @@ Swagger openapi defination see [openapi.json](docs/openapi.json). You can import it in [Swagger-UI](https://swagger.io/tools/swagger-ui/) editor. +All the generation api support for response in PNG bytes directly when request's 'Accept' header is 'image/png'. + +All the generation api support async process by pass parameter `async_process`` to true. And then use query job api to retrieve progress and generation results. + #### Text to Image > POST /v1/generation/text-to-image Alternative api for the normal image generation of Fooocus Gradio interface. -Add support for response in PNG bytes directly when request's 'Accept' header is 'image/png'. - #### Image Upscale or Variation > POST /v1/generation/image-upscale-vary @@ -91,3 +90,8 @@ Alternative api for 'Inpaint or Outpaint' tab of Fooocus Gradio interface. > POST /v1/generation/image-prompt Alternative api for 'Image Prompt' tab of Fooocus Gradio interface. + +#### Query Job +> GET /v1/generation/query-job + +Query async generation request results, return job progress and generation results. \ No newline at end of file diff --git a/docs/openapi.json b/docs/openapi.json index d5f3ed2..9f8fb5d 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -5,6 +5,22 @@ "version": "0.1.0" }, "paths": { + "/": { + "get": { + "summary": "Home", + "operationId": "home__get", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": {} + } + } + } + } + } + }, "/v1/generation/text-to-image": { "post": { "summary": "Text2Img Generation", @@ -54,10 +70,17 @@ "content": { "application/json": { "schema": { - "type": "array", - "items": { - "$ref": "#/components/schemas/GeneratedImageBase64" - }, + "anyOf": [ + { + "type": "array", + "items": { + "$ref": "#/components/schemas/GeneratedImageBase64" + } + }, + { + "$ref": "#/components/schemas/AsyncJobResponse" + } + ], "title": "Response Text2Img Generation V1 Generation Text To Image Post" }, "example": [ @@ -68,6 +91,12 @@ } ] }, + "application/json async": { + "example": { + "job_id": 1, + "job_type": "Text to Image" + } + }, "image/png": { "example": "PNG bytes, what did you expect?" } @@ -135,10 +164,17 @@ "content": { "application/json": { "schema": { - "type": "array", - "items": { - "$ref": "#/components/schemas/GeneratedImageBase64" - }, + "anyOf": [ + { + "type": "array", + "items": { + "$ref": "#/components/schemas/GeneratedImageBase64" + } + }, + { + "$ref": "#/components/schemas/AsyncJobResponse" + } + ], "title": "Response Img Upscale Or Vary V1 Generation Image Upscale Vary Post" }, "example": [ @@ -149,6 +185,12 @@ } ] }, + "application/json async": { + "example": { + "job_id": 1, + "job_type": "Text to Image" + } + }, "image/png": { "example": "PNG bytes, what did you expect?" } @@ -216,10 +258,17 @@ "content": { "application/json": { "schema": { - "type": "array", - "items": { - "$ref": "#/components/schemas/GeneratedImageBase64" - }, + "anyOf": [ + { + "type": "array", + "items": { + "$ref": "#/components/schemas/GeneratedImageBase64" + } + }, + { + "$ref": "#/components/schemas/AsyncJobResponse" + } + ], "title": "Response Img Inpaint Or Outpaint V1 Generation Image Inpait Outpaint Post" }, "example": [ @@ -230,6 +279,12 @@ } ] }, + "application/json async": { + "example": { + "job_id": 1, + "job_type": "Text to Image" + } + }, "image/png": { "example": "PNG bytes, what did you expect?" } @@ -301,10 +356,17 @@ "content": { "application/json": { "schema": { - "type": "array", - "items": { - "$ref": "#/components/schemas/GeneratedImageBase64" - }, + "anyOf": [ + { + "type": "array", + "items": { + "$ref": "#/components/schemas/GeneratedImageBase64" + } + }, + { + "$ref": "#/components/schemas/AsyncJobResponse" + } + ], "title": "Response Img Prompt V1 Generation Image Prompt Post" }, "example": [ @@ -315,6 +377,12 @@ } ] }, + "application/json async": { + "example": { + "job_id": 1, + "job_type": "Text to Image" + } + }, "image/png": { "example": "PNG bytes, what did you expect?" } @@ -332,6 +400,46 @@ } } } + }, + "/v1/generation/query-job": { + "get": { + "summary": "Query Job", + "description": "Query async generation job", + "operationId": "query_job_v1_generation_query_job_get", + "parameters": [ + { + "name": "job_id", + "in": "query", + "required": true, + "schema": { + "type": "integer", + "title": "Job Id" + } + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AsyncJobResponse" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } } }, "components": { @@ -368,6 +476,69 @@ ], "title": "AspectRatio" }, + "AsyncJobResponse": { + "properties": { + "job_id": { + "type": "integer", + "title": "Job Id" + }, + "job_type": { + "$ref": "#/components/schemas/TaskType" + }, + "job_stage": { + "$ref": "#/components/schemas/AsyncJobStage" + }, + "job_progess": { + "type": "integer", + "title": "Job Progess" + }, + "job_status": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Job Status" + }, + "job_result": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/GeneratedImageBase64" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Job Result" + } + }, + "type": "object", + "required": [ + "job_id", + "job_type", + "job_stage", + "job_progess", + "job_status", + "job_result" + ], + "title": "AsyncJobResponse" + }, + "AsyncJobStage": { + "type": "string", + "enum": [ + "WAITING", + "RUNNING", + "SUCCESS", + "ERROR" + ], + "title": "AsyncJobStage" + }, "Body_img_inpaint_or_outpaint_v1_generation_image_inpait_outpaint_post": { "properties": { "input_image": { @@ -557,6 +728,12 @@ "minimum": -2, "title": "W5", "default": 0.5 + }, + "async_process": { + "type": "boolean", + "title": "Async Process", + "description": "Set to true will run async and return job info for retrieve generataion result later", + "default": false } }, "type": "object", @@ -585,7 +762,7 @@ } ], "title": "Cn Stop1", - "description": "Stop at for image prompt" + "description": "Stop at for image prompt, None for default value" }, "cn_weight1": { "anyOf": [ @@ -628,7 +805,7 @@ } ], "title": "Cn Stop2", - "description": "Stop at for image prompt" + "description": "Stop at for image prompt, None for default value" }, "cn_weight2": { "anyOf": [ @@ -671,7 +848,7 @@ } ], "title": "Cn Stop3", - "description": "Stop at for image prompt" + "description": "Stop at for image prompt, None for default value" }, "cn_weight3": { "anyOf": [ @@ -714,7 +891,7 @@ } ], "title": "Cn Stop4", - "description": "Stop at for image prompt" + "description": "Stop at for image prompt, None for default value" }, "cn_weight4": { "anyOf": [ @@ -905,6 +1082,12 @@ "minimum": -2, "title": "W5", "default": 0.5 + }, + "async_process": { + "type": "boolean", + "title": "Async Process", + "description": "Set to true will run async and return job info for retrieve generataion result later", + "default": false } }, "type": "object", @@ -1087,6 +1270,12 @@ "minimum": -2, "title": "W5", "default": 0.5 + }, + "async_process": { + "type": "boolean", + "title": "Async Process", + "description": "Set to true will run async and return job info for retrieve generataion result later", + "default": false } }, "type": "object", @@ -1105,196 +1294,6 @@ ], "title": "ControlNetType" }, - "FooocusStyle": { - "type": "string", - "enum": [ - "Fooocus V2", - "Default (Slightly Cinematic)", - "SAI 3D Model", - "SAI Analog Film", - "SAI Anime", - "SAI Cinematic", - "SAI Comic Book", - "SAI Craft Clay", - "SAI Digital Art", - "SAI Enhance", - "SAI Fantasy Art", - "SAI Isometric", - "SAI Line Art", - "SAI Lowpoly", - "SAI Neonpunk", - "SAI Prigami", - "SAI Photographic", - "SAI Pixel Art", - "SAI Texture", - "Ads Advertising", - "Ads Automotive", - "Ads Corporate", - "Ads Fashion Editorial", - "Ads Food Photography", - "Ads Luxury", - "Ads Real Estate", - "Ads Retail", - "Artstyle Abstract", - "Artstyle Abstract Expressionism", - "Artstyle Art Deco", - "Artstyle Art Nouveau", - "Artstyle Constructivist", - "Artstyle Cubist", - "Artstyle Expressionist", - "Artstyle Graffiti", - "Artstyle Hyperrealism", - "Artstyle Impressionist", - "Artstyle Pointillism", - "Artstyle Pop Art", - "Artstyle Psychedelic", - "Artstyle Renaissance", - "Artstyle Steampunk", - "Artstyle Surrealist", - "Artstyle Typography", - "Artstyle Watercolor", - "Futuristic Biomechanical", - "Futuristic Biomechanical Cyberpunk", - "Futuristic Cybernetic", - "Futuristic Cybernetic Robot", - "Futuristic Cyberpunk Cityscape", - "Futuristic Futuristic", - "Futuristic Retro Cyberpunk", - "Futuristic Retro Futurism", - "Futuristic Sci Fi", - "Futuristic Vaporwave", - "Game Bubble Bobble", - "Game Cyberpunk Game", - "Game Fighting Game", - "Game Gta", - "Game Mario", - "Game Minecraft", - "Game Pokemon", - "Game Retro Arcade", - "Game Retro Game", - "Game Rpg Fantasy Game", - "Game Strategy Game", - "Game Streetfighter", - "Game Zelda", - "Misc Architectural", - "Misc Disco", - "Misc Dreamscape", - "Misc Dystopian", - "Misc Fairy Tale", - "Misc Gothic", - "Misc Grunge", - "Misc Horror", - "Misc Kawaii", - "Misc Lovecraftian", - "Misc Macabre", - "Misc Manga", - "Misc Metropolis", - "Misc Minimalist", - "Misc Monochrome", - "Misc Nautical", - "Misc Space", - "Misc Stained Glass", - "Misc Techwear Fashion", - "Misc Tribal", - "Misc Zentangle", - "Papercraft Collage", - "Papercraft Flat Papercut", - "Papercraft Kirigami", - "Papercraft Paper Mache", - "Papercraft Paper Quilling", - "Papercraft Papercut Collage", - "Papercraft Papercut Shadow Box", - "Papercraft Stacked Papercut", - "Papercraft Thick Layered Papercut", - "Photo Alien", - "Photo Film Noir", - "Photo Hdr", - "Photo Long Exposure", - "Photo Neon Noir", - "Photo Silhouette", - "Photo Tilt Shift", - "Cinematic Diva", - "Abstract Expressionism", - "Academia", - "Action Figure", - "Adorable 3D Character", - "Adorable Kawaii", - "Art Deco", - "Art Nouveau", - "Astral Aura", - "Avant Garde", - "Baroque", - "Bauhaus Style Poster", - "Blueprint Schematic Drawing", - "Caricature", - "Cel Shaded Art", - "Character Design Sheet", - "Classicism Art", - "Color Field Painting", - "Colored Pencil Art", - "Conceptual Art", - "Constructivism", - "Cubism", - "Dadaism", - "Dark Fantasy", - "Dark Moody Atmosphere", - "DMT Art Style", - "Doodle Art", - "Double Exposure", - "Dripping Paint Splatter Art", - "Expressionism", - "Faded Polaroid Photo", - "Fauvism", - "Flat 2D Art", - "Fortnite Art Style", - "Futurism", - "Glitchcore", - "Glo Fi", - "Googie Art Style", - "Graffiti Art", - "Harlem Renaissance Art", - "High Fashion", - "Idyllic", - "Impressionism", - "Infographic Drawing", - "Ink Dripping Drawing", - "Japanese Ink Drawing", - "Knolling Photography", - "Light Cheery Atmosphere", - "Logo Design", - "Luxurious Elegance", - "Macro Photography", - "Mandola Art", - "Marker Drawing", - "Medievalism", - "Minimalism", - "Neo Baroque", - "Neo Byzantine", - "Neo Futurism", - "Neo Impressionism", - "Neo Rococo", - "Neoclassicism", - "Op Art", - "Ornate And Intricate", - "Pencil Sketch Drawing", - "Pop Art 2", - "Rococo", - "Silhouette Art", - "Simple Vector Art", - "Sketchup", - "Steampunk 2", - "Surrealism", - "Suprematism", - "Terragen", - "Tranquil Relaxing Atmosphere", - "Sticker Designs", - "Vibrant Rim Light", - "Volumetric Lighting", - "Watercolor 2", - "Whimsical And Playful" - ], - "title": "FooocusStyle" - }, "GeneratedImageBase64": { "properties": { "base64": { @@ -1377,6 +1376,16 @@ ], "title": "PerfomanceSelection" }, + "TaskType": { + "type": "string", + "enum": [ + "Text to Image", + "Image Upscale or Variation", + "Image Inpaint or Outpaint", + "Image Prompt" + ], + "title": "TaskType" + }, "Text2ImgRequest": { "properties": { "prompt": { @@ -1391,7 +1400,7 @@ }, "style_selections": { "items": { - "$ref": "#/components/schemas/FooocusStyle" + "type": "string" }, "type": "array", "title": "Style Selections", @@ -1466,6 +1475,12 @@ "weight": 0.5 } ] + }, + "async_process": { + "type": "boolean", + "title": "Async Process", + "description": "Set to true will run async and return job info for retrieve generataion result later", + "default": false } }, "type": "object", diff --git a/fooocus_api_version.py b/fooocus_api_version.py index 68f9d21..2f8d1fc 100644 --- a/fooocus_api_version.py +++ b/fooocus_api_version.py @@ -1 +1 @@ -version = '0.1.20' \ No newline at end of file +version = '0.2.1' \ No newline at end of file diff --git a/fooocusapi/api.py b/fooocusapi/api.py index b027adb..586e980 100644 --- a/fooocusapi/api.py +++ b/fooocusapi/api.py @@ -1,15 +1,17 @@ from typing import List, Optional -from fastapi import Depends, FastAPI, Header, Query, UploadFile +from fastapi import Depends, FastAPI, Header, Query, Response, UploadFile from fastapi.params import File import uvicorn 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 +from fooocusapi.models import AsyncJobResponse, GeneratedImageBase64, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest +from fooocusapi.parameters import GenerationFinishReason, ImageGenerationResult +from fooocusapi.task_queue import TaskType +from fooocusapi.worker import process_generate, task_queue +from concurrent.futures import ThreadPoolExecutor app = FastAPI() -task_queue = TaskQueue() +work_executor = ThreadPoolExecutor(max_workers=12, thread_name_prefix="worker_") img_generate_responses = { "200": { @@ -22,6 +24,12 @@ "finish_reason": "SUCCESS" }] }, + "application/json async": { + "example": { + "job_id": 1, + "job_type": "Text to Image" + } + }, "image/png": { "example": "PNG bytes, what did you expect?" } @@ -29,12 +37,36 @@ } } +def call_worker(req: Text2ImgRequest, accept: str): + task_type = TaskType.text_2_img + if isinstance(req, ImgUpscaleOrVaryRequest): + task_type = TaskType.img_uov + elif isinstance(req, ImgInpaintOrOutpaintRequest): + task_type = TaskType.img_inpaint_outpaint + elif isinstance(req, ImgPromptRequest): + task_type = TaskType.img_prompt + + params = req_to_params(req) + queue_task = task_queue.add_task(task_type, {'params': params.__dict__, 'accept': accept}) + + if queue_task is None: + print("[Task Queue] The task queue has reached limit") + results = [ImageGenerationResult(im=None, seed=0, + finish_reason=GenerationFinishReason.queue_is_full)] + elif req.async_process: + work_executor.submit(process_generate, queue_task, params) + results = queue_task + else: + results = process_generate(queue_task, params) + + return results + @app.get("/") def home(): - return "Swagger-UI to: /docs" + return Response(content='Swagger-UI to: /docs', media_type="text/html") -@app.post("/v1/generation/text-to-image", response_model=List[GeneratedImageBase64], responses=img_generate_responses) +@app.post("/v1/generation/text-to-image", response_model=List[GeneratedImageBase64] | AsyncJobResponse, responses=img_generate_responses) def text2img_generation(req: Text2ImgRequest, accept: str = Header(None), accept_query: str | None = Query(None, alias='accept', description="Parameter to overvide 'Accept' header, 'image/png' for output bytes")): if accept_query is not None and len(accept_query) > 0: @@ -47,11 +79,11 @@ def text2img_generation(req: Text2ImgRequest, accept: str = Header(None), else: streaming_output = False - results = process_generate(req_to_params(req)) + results = call_worker(req, accept) return generation_output(results, streaming_output) -@app.post("/v1/generation/image-upscale-vary", response_model=List[GeneratedImageBase64], responses=img_generate_responses) +@app.post("/v1/generation/image-upscale-vary", response_model=List[GeneratedImageBase64] | AsyncJobResponse, responses=img_generate_responses) 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 overvide 'Accept' header, 'image/png' for output bytes")): @@ -65,11 +97,11 @@ def img_upscale_or_vary(input_image: UploadFile, req: ImgUpscaleOrVaryRequest = else: streaming_output = False - results = process_generate(req_to_params(req)) + results = call_worker(req, accept) return generation_output(results, streaming_output) -@app.post("/v1/generation/image-inpait-outpaint", response_model=List[GeneratedImageBase64], responses=img_generate_responses) +@app.post("/v1/generation/image-inpait-outpaint", response_model=List[GeneratedImageBase64] | AsyncJobResponse, responses=img_generate_responses) 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 overvide 'Accept' header, 'image/png' for output bytes")): @@ -83,11 +115,11 @@ def img_inpaint_or_outpaint(input_image: UploadFile, req: ImgInpaintOrOutpaintRe else: streaming_output = False - results = process_generate(req_to_params(req)) + results = call_worker(req, accept) return generation_output(results, streaming_output) -@app.post("/v1/generation/image-prompt", response_model=List[GeneratedImageBase64], responses=img_generate_responses) +@app.post("/v1/generation/image-prompt", response_model=List[GeneratedImageBase64] | AsyncJobResponse, responses=img_generate_responses) def img_prompt(cn_img1: Optional[UploadFile] = File(None), req: ImgPromptRequest = Depends(ImgPromptRequest.as_form), accept: str = Header(None), @@ -102,10 +134,19 @@ def img_prompt(cn_img1: Optional[UploadFile] = File(None), else: streaming_output = False - results = process_generate(req_to_params(req)) + results = call_worker(req, accept) return generation_output(results, streaming_output) +@app.get("/v1/generation/query-job", response_model=AsyncJobResponse, description="Query async generation job") +def query_job(job_id: int): + queue_task = task_queue.get_task(job_id, True) + if queue_task is None: + return Response(content="Job not found", status_code=404) + + return generation_output(queue_task, False) + + def start_app(args): uvicorn.run("fooocusapi.api:app", host=args.host, port=args.port, log_level=args.log_level) diff --git a/fooocusapi/api_utils.py b/fooocusapi/api_utils.py index d1cf2ad..969aa33 100644 --- a/fooocusapi/api_utils.py +++ b/fooocusapi/api_utils.py @@ -1,13 +1,14 @@ import base64 import io from io import BytesIO -from typing import List +from typing import List, Tuple import numpy as np from fastapi import Response, UploadFile from PIL import Image -from fooocusapi.models import GeneratedImageBase64, GenerationFinishReason, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest +from fooocusapi.models import AsyncJobResponse, AsyncJobStage, GeneratedImageBase64, GenerationFinishReason, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, Text2ImgRequest from fooocusapi.parameters import ImageGenerationParams, ImageGenerationResult +from fooocusapi.task_queue import QueueTask, TaskType import modules.flags as flags from modules.sdxl_styles import legal_style_names @@ -45,7 +46,8 @@ def read_input_image(input_image: UploadFile) -> np.ndarray: def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams: prompt = req.prompt negative_prompt = req.negative_prompt - style_selections = [s for s in req.style_selections if s in legal_style_names] + 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.value image_number = req.image_number @@ -102,7 +104,27 @@ def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams: ) -def generation_output(results: List[ImageGenerationResult], streaming_output: bool) -> Response | List[GeneratedImageBase64]: +def generation_output(results: QueueTask | List[ImageGenerationResult], streaming_output: bool) -> Response | List[GeneratedImageBase64] | AsyncJobResponse: + if isinstance(results, QueueTask): + task = results + job_stage = AsyncJobStage.running + job_result = None + if task.start_millis == 0: + job_stage = AsyncJobStage.waiting + if task.is_finished: + if task.finish_with_error: + job_stage = AsyncJobStage.error + else: + if task.task_result != None: + job_stage = AsyncJobStage.success + job_result = generation_output(task.task_result, False) + return AsyncJobResponse(job_id=task.seq, + job_type=task.type, + job_stage=job_stage, + job_progess=task.finish_progess, + job_status=task.task_status, + job_result=job_result) + if streaming_output: if len(results) == 0 or results[0].finish_reason != GenerationFinishReason.success: return Response(status_code=500) diff --git a/fooocusapi/models.py b/fooocusapi/models.py index 0a9b97e..cd3f74a 100644 --- a/fooocusapi/models.py +++ b/fooocusapi/models.py @@ -7,6 +7,7 @@ from pydantic_core import InitErrorDetails from fooocusapi.parameters import GenerationFinishReason, defualt_styles +from fooocusapi.task_queue import TaskType import modules.flags as flags @@ -97,6 +98,7 @@ class Text2ImgRequest(BaseModel): refiner_model_name: str = 'sd_xl_refiner_1.0_0.9vae.safetensors' loras: List[Lora] = Field(default=[ Lora(model_name='sd_xl_offset_example-lora_1.0.safetensors', weight=0.5)]) + async_process: bool = Field(default=False, description="Set to true will run async and return job info for retrieve generataion result later") class ImgUpscaleOrVaryRequest(Text2ImgRequest): @@ -133,12 +135,14 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for upsa w4: float = Form(default=0.5, ge=-2, le=2), l5: str | None = Form(None), w5: float = Form(default=0.5, ge=-2, le=2), + async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generataion result later"), ): style_selection_arr: List[str] = [] for part in style_selections: if len(part) > 0: for s in part.split(','): - style_selection_arr.append(s) + style = s.strip() + style_selection_arr.append(style) loras: List[Lora] = [] lora_config = [(l1, w1), (l2, w2), (l3, w3), (l4, w4), (l5, w5)] @@ -151,7 +155,7 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for upsa 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) + loras=loras, async_process=async_process) class ImgInpaintOrOutpaintRequest(Text2ImgRequest): @@ -192,6 +196,7 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for inpa w4: float = Form(default=0.5, ge=-2, le=2), l5: str | None = Form(None), w5: float = Form(default=0.5, ge=-2, le=2), + async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generataion result later"), ): if isinstance(input_mask, File): @@ -213,7 +218,8 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for inpa for part in style_selections: if len(part) > 0: for s in part.split(','): - style_selection_arr.append(s) + style = s.strip() + style_selection_arr.append(style) loras: List[Lora] = [] lora_config = [(l1, w1), (l2, w2), (l3, w3), (l4, w4), (l5, w5)] @@ -226,7 +232,7 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for inpa 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) + loras=loras, async_process=async_process) class ImgPromptRequest(Text2ImgRequest): @@ -291,6 +297,7 @@ def as_form(cls, cn_img1: UploadFile = Form(File(None), description="Input image w4: float = Form(default=0.5, ge=-2, le=2), l5: str | None = Form(None), w5: float = Form(default=0.5, ge=-2, le=2), + async_process: bool = Form(default=False, description="Set to true will run async and return job info for retrieve generataion result later"), ): if isinstance(cn_img1, File): cn_img1 = None @@ -317,7 +324,8 @@ def as_form(cls, cn_img1: UploadFile = Form(File(None), description="Input image for part in style_selections: if len(part) > 0: for s in part.split(','): - style_selection_arr.append(s) + style = s.strip() + style_selection_arr.append(style) loras: List[Lora] = [] lora_config = [(l1, w1), (l2, w2), (l3, w3), (l4, w4), (l5, w5)] @@ -330,7 +338,7 @@ def as_form(cls, cn_img1: UploadFile = Form(File(None), description="Input image 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) + loras=loras, async_process=async_process) class GeneratedImageBase64(BaseModel): @@ -338,3 +346,19 @@ class GeneratedImageBase64(BaseModel): 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 AsyncJobStage(str, Enum): + waiting = 'WAITING' + running = 'RUNNING' + success = 'SUCCESS' + error = 'ERROR' + + +class AsyncJobResponse(BaseModel): + job_id: int + job_type: TaskType + job_stage: AsyncJobStage + job_progess: int + job_status: str | None + job_result: List[GeneratedImageBase64] | None \ No newline at end of file diff --git a/fooocusapi/parameters.py b/fooocusapi/parameters.py index b8b0f2f..a083855 100644 --- a/fooocusapi/parameters.py +++ b/fooocusapi/parameters.py @@ -1,7 +1,9 @@ from enum import Enum -from typing import BinaryIO, Dict, List, Tuple +from typing import Dict, List, Tuple import numpy as np +from fooocusapi.task_queue import TaskType + inpaint_model_version = 'v1' @@ -9,222 +11,6 @@ defualt_styles = ['Fooocus V2', 'Default (Slightly Cinematic)'] -fooocus_styles = [ - 'Fooocus V2', - 'Default (Slightly Cinematic)', - 'Fooocus Masterpiece', - 'Fooocus Photograph', - 'Fooocus Negative', - 'SAI 3D Model', - 'SAI Analog Film', - 'SAI Anime', - 'SAI Cinematic', - 'SAI Comic Book', - 'SAI Craft Clay', - 'SAI Digital Art', - 'SAI Enhance', - 'SAI Fantasy Art', - 'SAI Isometric', - 'SAI Line Art', - 'SAI Lowpoly', - 'SAI Neonpunk', - 'SAI Origami', - 'SAI Photographic', - 'SAI Pixel Art', - 'SAI Texture', - 'MRE Cinematic Dynamic', - 'MRE Spontaneous Picture', - 'MRE Artistic Vision', - 'MRE Dark Dream', - 'MRE Gloomy Art', - 'MRE Bad Dream', - 'MRE Underground', - 'MRE Surreal Painting', - 'MRE Dynamic Illustration', - 'MRE Undead Art', - 'MRE Elemental Art', - 'MRE Space Art', - 'MRE Ancient Illustration', - 'MRE Brave Art', - 'MRE Heroic Fantasy', - 'MRE Dark Cyberpunk', - 'MRE Lyrical Geometry', - 'MRE Sumi E Symbolic', - 'MRE Sumi E Detailed', - 'MRE Manga', - 'MRE Anime', - 'MRE Comic', - 'Ads Advertising', - 'Ads Automotive', - 'Ads Corporate', - 'Ads Fashion Editorial', - 'Ads Food Photography', - 'Ads Gourmet Food Photography', - 'Ads Luxury', - 'Ads Real Estate', - 'Ads Retail', - 'Artstyle Abstract', - 'Artstyle Abstract Expressionism', - 'Artstyle Art Deco', - 'Artstyle Art Nouveau', - 'Artstyle Constructivist', - 'Artstyle Cubist', - 'Artstyle Expressionist', - 'Artstyle Graffiti', - 'Artstyle Hyperrealism', - 'Artstyle Impressionist', - 'Artstyle Pointillism', - 'Artstyle Pop Art', - 'Artstyle Psychedelic', - 'Artstyle Renaissance', - 'Artstyle Steampunk', - 'Artstyle Surrealist', - 'Artstyle Typography', - 'Artstyle Watercolor', - 'Futuristic Biomechanical', - 'Futuristic Biomechanical Cyberpunk', - 'Futuristic Cybernetic', - 'Futuristic Cybernetic Robot', - 'Futuristic Cyberpunk Cityscape', - 'Futuristic Futuristic', - 'Futuristic Retro Cyberpunk', - 'Futuristic Retro Futurism', - 'Futuristic Sci Fi', - 'Futuristic Vaporwave', - 'Game Bubble Bobble', - 'Game Cyberpunk Game', - 'Game Fighting Game', - 'Game Gta', - 'Game Mario', - 'Game Minecraft', - 'Game Pokemon', - 'Game Retro Arcade', - 'Game Retro Game', - 'Game Rpg Fantasy Game', - 'Game Strategy Game', - 'Game Streetfighter', - 'Game Zelda', - 'Misc Architectural', - 'Misc Disco', - 'Misc Dreamscape', - 'Misc Dystopian', - 'Misc Fairy Tale', - 'Misc Gothic', - 'Misc Grunge', - 'Misc Horror', - 'Misc Kawaii', - 'Misc Lovecraftian', - 'Misc Macabre', - 'Misc Manga', - 'Misc Metropolis', - 'Misc Minimalist', - 'Misc Monochrome', - 'Misc Nautical', - 'Misc Space', - 'Misc Stained Glass', - 'Misc Techwear Fashion', - 'Misc Tribal', - 'Misc Zentangle', - 'Papercraft Collage', - 'Papercraft Flat Papercut', - 'Papercraft Kirigami', - 'Papercraft Paper Mache', - 'Papercraft Paper Quilling', - 'Papercraft Papercut Collage', - 'Papercraft Papercut Shadow Box', - 'Papercraft Stacked Papercut', - 'Papercraft Thick Layered Papercut', - 'Photo Alien', - 'Photo Film Noir', - 'Photo Glamour', - 'Photo Hdr', - 'Photo Iphone Photographic', - 'Photo Long Exposure', - 'Photo Neon Noir', - 'Photo Silhouette', - 'Photo Tilt Shift', - 'Cinematic Diva', - 'Abstract Expressionism', - 'Academia', - 'Action Figure', - 'Adorable 3D Character', - 'Adorable Kawaii', - 'Art Deco', - 'Art Nouveau', - 'Astral Aura', - 'Avant Garde', - 'Baroque', - 'Bauhaus Style Poster', - 'Blueprint Schematic Drawing', - 'Caricature', - 'Cel Shaded Art', - 'Character Design Sheet', - 'Classicism Art', - 'Color Field Painting', - 'Colored Pencil Art', - 'Conceptual Art', - 'Constructivism', - 'Cubism', - 'Dadaism', - 'Dark Fantasy', - 'Dark Moody Atmosphere', - 'Dmt Art Style', - 'Doodle Art', - 'Double Exposure', - 'Dripping Paint Splatter Art', - 'Expressionism', - 'Faded Polaroid Photo', - 'Fauvism', - 'Flat 2d Art', - 'Fortnite Art Style', - 'Futurism', - 'Glitchcore', - 'Glo Fi', - 'Googie Art Style', - 'Graffiti Art', - 'Harlem Renaissance Art', - 'High Fashion', - 'Idyllic', - 'Impressionism', - 'Infographic Drawing', - 'Ink Dripping Drawing', - 'Japanese Ink Drawing', - 'Knolling Photography', - 'Light Cheery Atmosphere', - 'Logo Design', - 'Luxurious Elegance', - 'Macro Photography', - 'Mandola Art', - 'Marker Drawing', - 'Medievalism', - 'Minimalism', - 'Neo Baroque', - 'Neo Byzantine', - 'Neo Futurism', - 'Neo Impressionism', - 'Neo Rococo', - 'Neoclassicism', - 'Op Art', - 'Ornate And Intricate', - 'Pencil Sketch Drawing', - 'Pop Art 2', - 'Rococo', - 'Silhouette Art', - 'Simple Vector Art', - 'Sketchup', - 'Steampunk 2', - 'Surrealism', - 'Suprematism', - 'Terragen', - 'Tranquil Relaxing Atmosphere', - 'Sticker Designs', - 'Vibrant Rim Light', - 'Volumetric Lighting', - 'Watercolor 2', - 'Whimsical And Playful' -] - - aspect_ratios = [ '704×1408', '704×1344', diff --git a/fooocusapi/task_queue.py b/fooocusapi/task_queue.py index 86ade16..f21130e 100644 --- a/fooocusapi/task_queue.py +++ b/fooocusapi/task_queue.py @@ -4,15 +4,21 @@ class TaskType(str, Enum): - text2img = 'text2img' + text_2_img = 'Text to Image' + img_uov = 'Image Upscale or Variation' + img_inpaint_outpaint = 'Image Inpaint or Outpaint' + img_prompt = 'Image Prompt' class QueueTask(object): is_finished: bool = False + finish_progess: int = 0 start_millis: int = 0 finish_millis: int = 0 finish_with_error: bool = False + task_status: str | None = None task_result: any = None + error_message: str | None = None def __init__(self, seq: int, type: TaskType, req_param: dict, in_queue_millis: int): self.seq = seq @@ -20,14 +26,29 @@ def __init__(self, seq: int, type: TaskType, req_param: dict, in_queue_millis: i self.req_param = req_param self.in_queue_millis = in_queue_millis + def set_progress(self, progress: int, status: str | None): + if progress > 100: + progress = 100 + self.finish_progess = progress + self.task_status = status + + def set_result(self, task_result: any, finish_with_error: bool, error_message: str | None = None): + if not finish_with_error: + self.finish_progess = 100 + self.task_status = 'Finished' + self.task_result = task_result + self.finish_with_error = finish_with_error + self.error_message = error_message + class TaskQueue(object): queue: List[QueueTask] = [] history: List[QueueTask] = [] last_seq = 0 - def __init__(self, queue_size: int = 3): + def __init__(self, queue_size: int, hisotry_size: int): self.queue_size = queue_size + self.hisotry_size = hisotry_size def add_task(self, type: TaskType, req_param: dict) -> QueueTask | None: """ @@ -40,8 +61,8 @@ def add_task(self, type: TaskType, req_param: dict) -> QueueTask | None: task = QueueTask(seq=self.last_seq+1, type=type, req_param=req_param, in_queue_millis=int(round(time.time() * 1000))) self.queue.append(task) - self.last_seq += task.seq - return task.seq + self.last_seq = task.seq + return task def get_task(self, seq: int, include_history: bool = False) -> QueueTask | None: for task in self.queue: @@ -67,14 +88,17 @@ def start_task(self, seq: int): if task is not None: task.start_millis = int(round(time.time() * 1000)) - def finish_task(self, seq: int, task_result: any, finish_with_error: bool): + def finish_task(self, seq: int): task = self.get_task(seq) if task is not None: task.is_finished = True task.finish_millis = int(round(time.time() * 1000)) - task.finish_with_error = finish_with_error - task.task_result = task_result # Move task to history self.queue.remove(task) self.history.append(task) + + # Clean history + if len(self.history) > self.hisotry_size: + removed_task = self.history.pop(0) + print(f"Clean task history, remove task: {removed_task.seq}") diff --git a/fooocusapi/worker.py b/fooocusapi/worker.py index 6891667..b8c322d 100644 --- a/fooocusapi/worker.py +++ b/fooocusapi/worker.py @@ -5,15 +5,15 @@ import torch from typing import List from fooocusapi.parameters import inpaint_model_version, GenerationFinishReason, ImageGenerationParams, ImageGenerationResult -from fooocusapi.task_queue import TaskQueue, TaskType +from fooocusapi.task_queue import QueueTask, TaskQueue save_log = True -task_queue = TaskQueue() +task_queue = TaskQueue(queue_size=3, hisotry_size=6) @torch.no_grad() @torch.inference_mode() -def process_generate(params: ImageGenerationParams) -> List[ImageGenerationResult]: +def process_generate(queue_task: QueueTask, params: ImageGenerationParams) -> List[ImageGenerationResult]: import modules.default_pipeline as pipeline import modules.patch as patch import modules.flags as flags @@ -35,42 +35,39 @@ def process_generate(params: ImageGenerationParams) -> List[ImageGenerationResul def progressbar(number, text): print(f'[Fooocus] {text}') outputs.append(['preview', (number, text, None)]) + queue_task.set_progress(number, text) def make_results_from_outputs(): results: List[ImageGenerationResult] = [] for item in outputs: + seed = -1 if len(item) < 3 else item[2] if item[0] == 'results': for im in item[1]: if isinstance(im, np.ndarray): - results.append(ImageGenerationResult(im=im, seed=item[2], finish_reason=GenerationFinishReason.success)) - return results - - task_seq = task_queue.add_task(TaskType.text2img, { - 'body': params.__dict__}) - if task_seq is None: - print("[Task Queue] The task queue has reached limit") - results = [ImageGenerationResult(im=None, seed=0, - finish_reason=GenerationFinishReason.queue_is_full)] + results.append(ImageGenerationResult(im=im, seed=seed, finish_reason=GenerationFinishReason.success)) + queue_task.set_result(results, False) + task_queue.finish_task(queue_task.seq) + print(f"[Task Queue] Finish task, seq={queue_task.seq}") return results try: waiting_sleep_steps: int = 0 waiting_start_time = time.perf_counter() - while not task_queue.is_task_ready_to_start(task_seq): + while not task_queue.is_task_ready_to_start(queue_task.seq): if waiting_sleep_steps == 0: print( - f"[Task Queue] Waiting for task queue become free, seq={task_seq}") + f"[Task Queue] Waiting for task queue become free, seq={queue_task.seq}") delay = 0.1 time.sleep(delay) waiting_sleep_steps += 1 if waiting_sleep_steps % int(10 / delay) == 0: waiting_time = time.perf_counter() - waiting_start_time print( - f"[Task Queue] Already waiting for {waiting_time}S, seq={task_seq}") + f"[Task Queue] Already waiting for {waiting_time}S, seq={queue_task.seq}") - print(f"[Task Queue] Task queue is free, start task, seq={task_seq}") + print(f"[Task Queue] Task queue is free, start task, seq={queue_task.seq}") - task_queue.start_task(task_seq) + task_queue.start_task(queue_task.seq) execution_start_time = time.perf_counter() @@ -401,10 +398,10 @@ def build_advanced_parameters(): if direct_return: d = [('Upscale (Fast)', '2x')] - log(uov_input_image, d, single_line_number=1) + if save_log: + log(uov_input_image, d, single_line_number=1) outputs.append(['results', [uov_input_image], -1 if len(tasks) == 0 else tasks[0]['task_seed']]) results = make_results_from_outputs() - task_queue.finish_task(task_seq, results, False) return results * image_number tiled = True @@ -456,9 +453,8 @@ def build_advanced_parameters(): pipeline.final_unet.model.diffusion_model.in_inpaint = True if advanced_parameters.debugging_cn_preprocessor: - outputs.append(['results', inpaint_worker.current_task.visualize_mask_processing()]) - results = [] - task_queue.finish_task(task_seq, results, False) + outputs.append(['results', inpaint_worker.current_task.visualize_mask_processing(), -1 if len(tasks) == 0 else tasks[0]['task_seed']]) + results = make_results_from_outputs() return results progressbar(13, 'VAE Inpaint encoding ...') @@ -506,7 +502,6 @@ def build_advanced_parameters(): if advanced_parameters.debugging_cn_preprocessor: outputs.append(['results', [cn_img], task['task_seed']]) results = make_results_from_outputs() - task_queue.finish_task(task_seq, results, False) return results for task in cn_tasks[flags.cn_cpds]: cn_img, cn_stop, cn_weight = task @@ -517,7 +512,6 @@ def build_advanced_parameters(): if advanced_parameters.debugging_cn_preprocessor: outputs.append(['results', [cn_img], task['task_seed']]) results = make_results_from_outputs() - task_queue.finish_task(task_seq, results, False) return results for task in cn_tasks[flags.cn_ip]: cn_img, cn_stop, cn_weight = task @@ -530,7 +524,6 @@ def build_advanced_parameters(): if advanced_parameters.debugging_cn_preprocessor: outputs.append(['results', [cn_img], task['task_seed']]) results = make_results_from_outputs() - task_queue.finish_task(task_seq, results, False) return results if len(cn_tasks[flags.cn_ip]) > 0: @@ -620,33 +613,31 @@ def callback(step, x0, x, total_steps, y): for n, w in loras_raw: if n != 'None': d.append((f'LoRA [{n}] weight', w)) - log(x, d, single_line_number=3) + if save_log: + log(x, d, single_line_number=3) # Fooocus async_worker.py code end 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(ImageGenerationResult( - im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.user_cancel)) - break except Exception as e: - print('Process failed:', e) + print('Process error:', e) results.append(ImageGenerationResult( im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.error)) + queue_task.set_result(results, True, str(e)) + break execution_time = time.perf_counter() - execution_start_time print(f'Generating and saving time: {execution_time:.2f} seconds') pipeline.prepare_text_encoder(async_call=True) - print(f"[Task Queue] Finish task, seq={task_seq}") - task_queue.finish_task(task_seq, results, False) - + if not queue_task.finish_with_error: + queue_task.set_result(results, False) + task_queue.finish_task(queue_task.seq) + print(f"[Task Queue] Finish task, seq={queue_task.seq}") return results except Exception as e: print('Worker error:', e) - print(f"[Task Queue] Finish task, seq={task_seq}") - task_queue.finish_task(task_seq, [], True) - raise e + queue_task.set_result([], True, str(e)) + return [] diff --git a/predict.py b/predict.py index 6aa9e47..1fed729 100644 --- a/predict.py +++ b/predict.py @@ -6,7 +6,7 @@ from typing import List from cog import BasePredictor, Input, Path -from fooocusapi.parameters import inpaint_model_version, GenerationFinishReason, ImageGenerationParams, fooocus_styles, aspect_ratios, uov_methods, outpaint_expansions +from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams, aspect_ratios, uov_methods, outpaint_expansions from fooocusapi.worker import process_generate import numpy as np from PIL import Image @@ -24,40 +24,68 @@ def setup(self) -> None: def predict( self, - prompt: str = Input(default='', description="Prompt for image generation"), - negative_prompt: str = Input(default='', description="Negtive prompt for image generation"), - style_selections: str = Input(default='Fooocus V2,Default (Slightly Cinematic)', description="Fooocus styles applied for image generation, seperated by comma"), - performance_selection: str = Input(default='Speed', description="Performance selection", choices=['Speed', 'Quality']), - aspect_ratios_selection: str = Input(default='1152×896', description="The generated image's size", choices=aspect_ratios), - image_number: int = Input(default=1, description="How many image to generate", ge=1, le=8), - image_seed: int = Input(default=-1, description="Seed to generate image, -1 for random"), + prompt: str = Input( + default='', description="Prompt for image generation"), + negative_prompt: str = Input( + default='', description="Negtive prompt for image generation"), + style_selections: str = Input(default='Fooocus V2,Default (Slightly Cinematic)', + description="Fooocus styles applied for image generation, seperated by comma"), + performance_selection: str = Input( + default='Speed', description="Performance selection", choices=['Speed', 'Quality']), + aspect_ratios_selection: str = Input( + default='1152×896', description="The generated image's size", choices=aspect_ratios), + image_number: int = Input( + default=1, description="How many image to generate", ge=1, le=8), + image_seed: int = Input( + default=-1, description="Seed to generate image, -1 for random"), sharpness: float = Input(default=2.0, ge=0.0, le=30.0), guidance_scale: float = Input(default=7.0, ge=1.0, le=30.0), - uov_input_image: Path = Input(default=None, description="Input image for upscale or variation, keep None for not upscale or variation"), + 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), - 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' seperated by comma"), - 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='Image Prompt', description="ControlNet type for image prompt", choices=['Image Prompt', 'PyraCanny', 'CPDS']), - 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='Image Prompt', description="ControlNet type for image prompt", choices=['Image Prompt', 'PyraCanny', 'CPDS']), - 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='Image Prompt', description="ControlNet type for image prompt", choices=['Image Prompt', 'PyraCanny', 'CPDS']), - 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='Image Prompt', description="ControlNet type for image prompt", choices=['Image Prompt', 'PyraCanny', 'CPDS']), + 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' seperated by comma"), + 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='Image Prompt', description="ControlNet type for image prompt", choices=[ + 'Image Prompt', 'PyraCanny', 'CPDS']), + 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='Image Prompt', description="ControlNet type for image prompt", choices=[ + 'Image Prompt', 'PyraCanny', 'CPDS']), + 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='Image Prompt', description="ControlNet type for image prompt", choices=[ + 'Image Prompt', 'PyraCanny', 'CPDS']), + 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='Image Prompt', description="ControlNet type for image prompt", choices=[ + 'Image Prompt', 'PyraCanny', 'CPDS']), ) -> List[Path]: """Run a single prediction on the model""" from modules.util import generate_temp_filename import modules.flags as flags + from modules.sdxl_styles import legal_style_names base_model_name = 'sd_xl_base_1.0_0.9vae.safetensors' refiner_model_name = 'sd_xl_refiner_1.0_0.9vae.safetensors' @@ -66,7 +94,7 @@ def predict( style_selections_arr = [] for s in style_selections.strip().split(','): style = s.strip() - if style in fooocus_styles: + if style in legal_style_names: style_selections_arr.append(style) if uov_input_image is not None: @@ -83,7 +111,7 @@ def predict( inpaint_input_mask = np.array(im) else: inpaint_input_mask = np.zeros(inpaint_input_image.shape) - + inpaint_input_image_dict = { 'image': inpaint_input_image, 'mask': inpaint_input_mask @@ -109,26 +137,25 @@ def predict( cn_weight = flags.default_parameters[cn_type][1] image_prompts.append((cn_img, cn_stop, cn_weight, cn_type)) - 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, - loras=loras, - uov_input_image=uov_input_image, - uov_method=uov_method, - outpaint_selections=outpaint_selections_arr, - inpaint_input_image=inpaint_input_image_dict, - image_prompts=image_prompts - ) - + 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, + loras=loras, + uov_input_image=uov_input_image, + uov_method=uov_method, + outpaint_selections=outpaint_selections_arr, + inpaint_input_image=inpaint_input_image_dict, + image_prompts=image_prompts + ) + print(f"[Predictor Predict] Params: {params.__dict__}") results = process_generate(params) @@ -137,7 +164,8 @@ def predict( 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) + 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))