From d3b5c90f4fda0e43163b1f2d9f28d3b5e0a234e5 Mon Sep 17 00:00:00 2001 From: Clemens Marschner Date: Wed, 8 Nov 2023 08:47:09 +0000 Subject: [PATCH 1/4] added server --- Makefile | 9 ++ .../detectron2/data/transforms/transform.py | 2 +- backend.py | 109 ++++++++++++++++++ celery_app.py | 62 ++++++++++ client.py | 78 +++++++++++++ environment.yaml | 5 + gradio_canny.py | 14 ++- server.py | 76 ++++++++++++ 8 files changed, 348 insertions(+), 7 deletions(-) create mode 100644 Makefile create mode 100644 backend.py create mode 100644 celery_app.py create mode 100644 client.py create mode 100644 server.py diff --git a/Makefile b/Makefile new file mode 100644 index 00000000..d3824a6e --- /dev/null +++ b/Makefile @@ -0,0 +1,9 @@ + +celery: + celery -A celery_app worker --loglevel=info --concurrency=1 --pool=solo + +server: + python server.py + +test: + python client.py img.png mask.png --prompt "a red drawer" diff --git a/annotator/oneformer/detectron2/data/transforms/transform.py b/annotator/oneformer/detectron2/data/transforms/transform.py index de44b991..46769a25 100644 --- a/annotator/oneformer/detectron2/data/transforms/transform.py +++ b/annotator/oneformer/detectron2/data/transforms/transform.py @@ -43,7 +43,7 @@ class ExtentTransform(Transform): See: https://pillow.readthedocs.io/en/latest/PIL.html#PIL.ImageTransform.ExtentTransform """ - def __init__(self, src_rect, output_size, interp=Image.LINEAR, fill=0): + def __init__(self, src_rect, output_size, interp=Image.BILINEAR, fill=0): """ Args: src_rect (x0, y0, x1, y1): src coordinates diff --git a/backend.py b/backend.py new file mode 100644 index 00000000..0f3fca18 --- /dev/null +++ b/backend.py @@ -0,0 +1,109 @@ +from share import * +import config + +import cv2 +import einops +import gradio as gr +import numpy as np +import torch +from typing import Dict, Any +import random + +from pytorch_lightning import seed_everything +from annotator.util import resize_image, HWC3 +from cldm.model import create_model, load_state_dict +from cldm.ddim_hacked import DDIMSampler + + +model = None +ddim_sampler = None + +def init_model(): + global model + global ddim_sampler + model_name = 'control_v11p_sd15_inpaint' + model = create_model(f'./models/{model_name}.yaml').cpu() + model.load_state_dict(load_state_dict('./models/v1-5-pruned.ckpt', location='cuda'), strict=False) + model.load_state_dict(load_state_dict(f'./models/{model_name}.pth', location='cuda'), strict=False) + model = model.cuda() + + ddim_sampler = DDIMSampler(model) + + +def process(input_image_and_mask, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, mask_blur, + update_state_fn): + assert prompt + with torch.no_grad(): + input_image = HWC3(input_image_and_mask['image']) + input_mask = input_image_and_mask['mask'] + + img_raw = resize_image(input_image, image_resolution).astype(np.float32) + H, W, C = img_raw.shape + + mask_pixel = cv2.resize(input_mask[:, :, 0], (W, H), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255.0 + mask_pixel = cv2.GaussianBlur(mask_pixel, (0, 0), mask_blur) + + mask_latent = cv2.resize(mask_pixel, (W // 8, H // 8), interpolation=cv2.INTER_AREA) + + detected_map = img_raw.copy() + detected_map[mask_pixel > 0.5] = - 255.0 + + control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0 + control = torch.stack([control for _ in range(num_samples)], dim=0) + control = einops.rearrange(control, 'b h w c -> b c h w').clone() + + mask = 1.0 - torch.from_numpy(mask_latent.copy()).float().cuda() + mask = torch.stack([mask for _ in range(num_samples)], dim=0) + mask = einops.rearrange(mask, 'b h w -> b 1 h w').clone() + + x0 = torch.from_numpy(img_raw.copy()).float().cuda() / 127.0 - 1.0 + x0 = torch.stack([x0 for _ in range(num_samples)], dim=0) + x0 = einops.rearrange(x0, 'b h w c -> b c h w').clone() + + mask_pixel_batched = mask_pixel[None, :, :, None] + img_pixel_batched = img_raw.copy()[None] + + if seed == -1: + seed = random.randint(0, 65535) + seed_everything(seed) + + if config.save_memory: + model.low_vram_shift(is_diffusing=False) + + prompts = [prompt + ', ' + a_prompt] * num_samples + cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning(prompts)]} + un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]} + shape = (4, H // 8, W // 8) + + if config.save_memory: + model.low_vram_shift(is_diffusing=False) + + ddim_sampler.make_schedule(ddim_steps, ddim_eta=eta, verbose=True) + x0 = model.get_first_stage_encoding(model.encode_first_stage(x0)) + + if config.save_memory: + model.low_vram_shift(is_diffusing=True) + + model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13) + # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01 + + def update_status_dict(i): + update_state_fn({"step": i, "num_steps": ddim_steps}) + + samples, intermediates = ddim_sampler.sample( + ddim_steps, num_samples, + shape, cond, verbose=False, eta=eta, + unconditional_guidance_scale=scale, + unconditional_conditioning=un_cond, x0=x0, mask=mask, + callback = update_status_dict + ) + + if config.save_memory: + model.low_vram_shift(is_diffusing=False) + + x_samples = model.decode_first_stage(samples) + x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().astype(np.float32) + x_samples = x_samples * mask_pixel_batched + img_pixel_batched * (1.0 - mask_pixel_batched) + + results = [x_samples[i].clip(0, 255).astype(np.uint8) for i in range(num_samples)] + return [detected_map.clip(0, 255).astype(np.uint8)] + results diff --git a/celery_app.py b/celery_app.py new file mode 100644 index 00000000..ad94f40f --- /dev/null +++ b/celery_app.py @@ -0,0 +1,62 @@ +from celery import Celery +from celery.signals import worker_process_init +import cv2 + +celery_app = Celery( + "worker", + broker="redis://localhost:6379/0", + backend="redis://localhost:6379/0" +) + +backend = None + +@worker_process_init.connect +def load_model_at_worker_init(*args, **kwargs): + import os + print("Current working directory:", os.getcwd()) + import sys + sys.path.append(".") + global backend + import backend + print('Loading model...') + backend.init_model() + print('Model loaded.') + +@celery_app.task(bind=True) +def handle_image_processing(self, image_filename, mask_filename, prompt): + image = cv2.imread(image_filename)[:,:,[2,1,0]] + mask = cv2.imread(mask_filename) + print("#### TASK #####") + try: + def update_state(state_dict): + self.update_state(state='PROGRESS', meta=state_dict) + + results: List[np.ndarray] = backend.process( + input_image_and_mask={ + "image": image, + "mask": mask + }, + prompt=prompt, + a_prompt="best quality", + n_prompt="lowres, bad anatomy, bad hands, cropped, worst quality", + num_samples=1, + image_resolution=512, + ddim_steps=20, + guess_mode=False, + strength=1, + scale=9, + seed=12345, + eta=1, + mask_blur=5, + update_state_fn=update_state, + ) + res_name = f"/tmp/tmp_{self.request.id}_result_1.png" + cv2.imwrite(res_name, results[0][:,:,[2,1,0]]) + return {"image_filename": res_name} + except Exception as e: + print(f"*** Exception in job {self.request.id}") + raise + # print(e) + finally: + # TODO clean-up + pass diff --git a/client.py b/client.py new file mode 100644 index 00000000..7025948e --- /dev/null +++ b/client.py @@ -0,0 +1,78 @@ +import requests +import time +import os + +# The endpoint for submitting the image processing request +SUBMIT_URL = 'http://localhost:8889/controlnet/' +# The endpoints for checking the status and getting the result +STATUS_URL = 'http://localhost:8889/jobs/status/{job_id}' +RESULT_URL = 'http://localhost:8889/jobs/result/{job_id}' + +def submit_image_processing(image_path, mask_path, coordinates, prompt=None): + files = { + 'image': ('image.png', open(image_path, 'rb'), 'image/png'), + 'mask': ('mask.png', open(mask_path, 'rb'), 'image/png'), + 'prompt': (None, prompt), # The first element is the filename, which is None in this case + } + data = { + } + response = requests.post(SUBMIT_URL, files=files) # , data=data) + return response.json() + +def check_status(job_id): + return requests.get(STATUS_URL.format(job_id=job_id)).json() + # check response.status_code + +def get_result(job_id, save_path): + response = requests.get(RESULT_URL.format(job_id=job_id)) + + if response.status_code == 200: + # Assuming the endpoint sends the file directly and there's no redirect + with open(save_path, 'wb') as f: + f.write(response.content) + print(f"Image successfully downloaded and saved to {save_path}") + else: + # Handle potential error (job not completed, result not available, etc.) + print("Status code: ", response.status_code) + print(response) + # error_info = response.json() + # print(f"Error: {error_info.get('error', 'Unknown error occurred')}") + +# Example usage +if __name__ == '__main__': + # Define the command-line argument parser + import argparse + parser = argparse.ArgumentParser(description='Submit an image processing job to the API') + parser.add_argument('image_path', type=str, help='Path to the RGB image file') + parser.add_argument('mask_path', type=str, help='Path to the boolean mask file') + # parser.add_argument('coordinates', type=str, help='JSON string of 3D coordinates and labels') + parser.add_argument('--prompt', type=str, default='', help='Optional prompt text') + + # Parse the command-line arguments + args = parser.parse_args() + + res = submit_image_processing( + args.image_path, + args.mask_path, + # args.coordinates, + args.prompt + ) + print("submitted with response", res) + job_id = res.get('job_id') + print(f"Job submitted. ID: {job_id}") + + if job_id: + # Wait and check the status until the job is completed + while True: + status_response = check_status(job_id) + job_status = status_response.get('status') + print(f"Job status: {status_response}") + if job_status == 'SUCCESS': + # Retrieve and save the result + save_path = 'output_image.png' + get_result(job_id, save_path) + break + elif job_status == 'FAILURE': + print(f"Error: {status_response.get('error')}") + break + time.sleep(2) diff --git a/environment.yaml b/environment.yaml index 4f87a9aa..c734d45f 100644 --- a/environment.yaml +++ b/environment.yaml @@ -9,6 +9,11 @@ dependencies: - pytorch=1.12.1 - torchvision=0.13.1 - numpy=1.23.1 + - jupyter + - fastapi + - uvicorn[standard] + - celery + - redis-py - pip: - gradio==3.16.2 - albumentations==1.3.0 diff --git a/gradio_canny.py b/gradio_canny.py index 217736c7..54ce79af 100644 --- a/gradio_canny.py +++ b/gradio_canny.py @@ -28,11 +28,13 @@ def process(det, input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, low_threshold, high_threshold): global preprocessor + print("Process") if det == 'Canny': if not isinstance(preprocessor, CannyDetector): preprocessor = CannyDetector() with torch.no_grad(): + print("image conversion") input_image = HWC3(input_image) if det == 'None': @@ -40,12 +42,12 @@ def process(det, input_image, prompt, a_prompt, n_prompt, num_samples, image_res else: detected_map = preprocessor(resize_image(input_image, detect_resolution), low_threshold, high_threshold) detected_map = HWC3(detected_map) - + print("resize") img = resize_image(input_image, image_resolution) H, W, C = img.shape detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - + print("to torch") control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0 control = torch.stack([control for _ in range(num_samples)], dim=0) control = einops.rearrange(control, 'b h w c -> b c h w').clone() @@ -66,7 +68,7 @@ def process(det, input_image, prompt, a_prompt, n_prompt, num_samples, image_res model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13) # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01 - + print("sampling") samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples, shape, cond, verbose=False, eta=eta, unconditional_guidance_scale=scale, @@ -74,10 +76,10 @@ def process(det, input_image, prompt, a_prompt, n_prompt, num_samples, image_res if config.save_memory: model.low_vram_shift(is_diffusing=False) - + print("decoding") x_samples = model.decode_first_stage(samples) x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8) - + print("returning samples") results = [x_samples[i] for i in range(num_samples)] return [detected_map] + results @@ -112,4 +114,4 @@ def process(det, input_image, prompt, a_prompt, n_prompt, num_samples, image_res run_button.click(fn=process, inputs=ips, outputs=[result_gallery]) -block.launch(server_name='0.0.0.0') +block.launch(server_name='0.0.0.0', share=False) diff --git a/server.py b/server.py new file mode 100644 index 00000000..e1328671 --- /dev/null +++ b/server.py @@ -0,0 +1,76 @@ +from fastapi import FastAPI, File, UploadFile, BackgroundTasks +from pydantic import BaseModel +from typing import List, Optional +from uuid import uuid4 +import shutil +import asyncio +from fastapi.responses import FileResponse +from celery_app import celery_app, handle_image_processing +from celery.result import AsyncResult + +app = FastAPI() + + +@app.post("/controlnet/") +async def controlnet(background_tasks: BackgroundTasks, + image: UploadFile, # = File(...), + mask: UploadFile, # = File(...), + # coordinates: List[Coordinate], + # prompt: str + ): + prompt = "red drawer" + print("Using prompt: ", prompt) + # Save files temporarily and handle them + image_filename = f"/tmp/temp_{uuid4()}_request_img.png" + mask_filename = f"/tmp/temp_{uuid4()}_request_mask.png" + print(f"saving to {image_filename} and {mask_filename}") + with open(image_filename, "wb") as buffer: + shutil.copyfileobj(image.file, buffer) + with open(mask_filename, "wb") as buffer: + shutil.copyfileobj(mask.file, buffer) + + # Create a job ID and start the background task + + task = handle_image_processing.delay(image_filename, mask_filename, prompt) + print(f"created job {task.id}") + return {"job_id": task.id} + + + +@app.get("/jobs/status/{job_id}") +async def get_job_status(job_id: str): + result = AsyncResult(job_id, app=celery_app) + if result.status == "FAILURE": + try: + res = result.get() + except Exception as e: + print("Exception: ", e) + return {"job_id": job_id, "status": result.status, "error": str(e)} + res_state = {"job_id": job_id, "status": result.status} + if result.status == "PROGRESS": + if result.info: + print(result.info) + res_state.update(result.info) + return res_state + +@app.get("/jobs/result/{job_id}") +async def get_result(job_id: str): + result = AsyncResult(job_id, app=celery_app) + try: + res = result.get() + except Exception as e: + print("Exception: ", e) + return {"job_id": job_id, "status": result.status, "error": str(e)} + + if result.status != "SUCCESS": + return {"job_id": job_id, "status": result.status, "error": "Result not available or job not completed."} + + image_path = res["image_filename"] + return FileResponse(image_path, media_type='image/png', filename=image_path.split("/")[-1]) + + +if __name__ == "__main__": + import uvicorn + uvicorn.run(app, host="0.0.0.0", port=8889) + + From 6b6492393c9a95b5bc56d45c0b9e7212c5cae030 Mon Sep 17 00:00:00 2001 From: Clemens Marschner Date: Sun, 12 Nov 2023 19:26:29 +0000 Subject: [PATCH 2/4] latest changes --- Makefile | 2 +- backend.py | 8 ++-- celery_app.py | 21 +++++++---- client.py | 96 +++++++++++++++++++++++++++++------------------ gradio_inpaint.py | 2 + server.py | 57 +++++++++++++++++----------- 6 files changed, 117 insertions(+), 69 deletions(-) diff --git a/Makefile b/Makefile index d3824a6e..a3f48105 100644 --- a/Makefile +++ b/Makefile @@ -6,4 +6,4 @@ server: python server.py test: - python client.py img.png mask.png --prompt "a red drawer" + python client.py img.png mask.png --prompt "a boatsteg made of glass overlooking a large lake" --seed 1 --num_images 1 --resolution 512 diff --git a/backend.py b/backend.py index 0f3fca18..12cb226f 100644 --- a/backend.py +++ b/backend.py @@ -40,7 +40,9 @@ def process(input_image_and_mask, prompt, a_prompt, n_prompt, num_samples, image img_raw = resize_image(input_image, image_resolution).astype(np.float32) H, W, C = img_raw.shape - mask_pixel = cv2.resize(input_mask[:, :, 0], (W, H), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255.0 + if (len(input_mask.shape) >= 3): + input_mask = input_mask[:, :, 0] + mask_pixel = cv2.resize(input_mask, (W, H), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255.0 mask_pixel = cv2.GaussianBlur(mask_pixel, (0, 0), mask_blur) mask_latent = cv2.resize(mask_pixel, (W // 8, H // 8), interpolation=cv2.INTER_AREA) @@ -88,11 +90,11 @@ def process(input_image_and_mask, prompt, a_prompt, n_prompt, num_samples, image # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01 def update_status_dict(i): - update_state_fn({"step": i, "num_steps": ddim_steps}) + update_state_fn({"step": i, "num_steps": ddim_steps + 1}) samples, intermediates = ddim_sampler.sample( ddim_steps, num_samples, - shape, cond, verbose=False, eta=eta, + shape, cond, verbose=True, eta=eta, unconditional_guidance_scale=scale, unconditional_conditioning=un_cond, x0=x0, mask=mask, callback = update_status_dict diff --git a/celery_app.py b/celery_app.py index ad94f40f..7a941365 100644 --- a/celery_app.py +++ b/celery_app.py @@ -23,9 +23,11 @@ def load_model_at_worker_init(*args, **kwargs): print('Model loaded.') @celery_app.task(bind=True) -def handle_image_processing(self, image_filename, mask_filename, prompt): +def handle_image_processing(self, image_filename, mask_filename, prompt, seed, num_images, resolution, num_steps): image = cv2.imread(image_filename)[:,:,[2,1,0]] mask = cv2.imread(mask_filename) + if mask.shape[0] != image.shape[0] or mask.shape[1] != image.shape[1]: + raise Exception(f"Expected image and mask to be of the same size, but got HxWxC {image.shape} vs. {mask.shape}") print("#### TASK #####") try: def update_state(state_dict): @@ -39,20 +41,23 @@ def update_state(state_dict): prompt=prompt, a_prompt="best quality", n_prompt="lowres, bad anatomy, bad hands, cropped, worst quality", - num_samples=1, - image_resolution=512, - ddim_steps=20, + num_samples=num_images, + image_resolution=resolution, + ddim_steps=num_steps, guess_mode=False, strength=1, scale=9, - seed=12345, + seed=seed, eta=1, mask_blur=5, update_state_fn=update_state, ) - res_name = f"/tmp/tmp_{self.request.id}_result_1.png" - cv2.imwrite(res_name, results[0][:,:,[2,1,0]]) - return {"image_filename": res_name} + images = [] + for i in range(len(results)): + res_name = f"/tmp/tmp_{self.request.id}_result_{i}.png" + cv2.imwrite(res_name, results[i][:,:,[2,1,0]]) + images.append(res_name) + return {"image_filenames": images, "num_images": len(results)} except Exception as e: print(f"*** Exception in job {self.request.id}") raise diff --git a/client.py b/client.py index 7025948e..3e65c9fa 100644 --- a/client.py +++ b/client.py @@ -1,42 +1,50 @@ import requests import time import os +import json +import sys # The endpoint for submitting the image processing request SUBMIT_URL = 'http://localhost:8889/controlnet/' # The endpoints for checking the status and getting the result STATUS_URL = 'http://localhost:8889/jobs/status/{job_id}' -RESULT_URL = 'http://localhost:8889/jobs/result/{job_id}' +RESULT_URL = 'http://localhost:8889/jobs/result/{job_id}/{img_id}' -def submit_image_processing(image_path, mask_path, coordinates, prompt=None): - files = { - 'image': ('image.png', open(image_path, 'rb'), 'image/png'), - 'mask': ('mask.png', open(mask_path, 'rb'), 'image/png'), - 'prompt': (None, prompt), # The first element is the filename, which is None in this case - } - data = { - } - response = requests.post(SUBMIT_URL, files=files) # , data=data) - return response.json() +def submit_image_processing(image_path, mask_path, prompt, resolution, num_images, seed, num_steps): + with open(image_path, 'rb') as r1, open(mask_path, 'rb') as r2: + files = { + 'image': ('image.png', r1, 'image/png'), + 'mask': ('mask.png', r2, 'image/png'), + } + params = { + 'prompt': prompt, + 'resolution': resolution, + 'num_images': num_images, + 'seed': seed, + 'num_steps': num_steps, + } + response = requests.post(SUBMIT_URL, params=params, files=files) # , data=data) + return response # .json() def check_status(job_id): return requests.get(STATUS_URL.format(job_id=job_id)).json() # check response.status_code def get_result(job_id, save_path): - response = requests.get(RESULT_URL.format(job_id=job_id)) - - if response.status_code == 200: - # Assuming the endpoint sends the file directly and there's no redirect - with open(save_path, 'wb') as f: - f.write(response.content) - print(f"Image successfully downloaded and saved to {save_path}") - else: - # Handle potential error (job not completed, result not available, etc.) - print("Status code: ", response.status_code) - print(response) - # error_info = response.json() - # print(f"Error: {error_info.get('error', 'Unknown error occurred')}") + status_obj = check_status(job_id) + for i in range(status_obj["num_images"]): + response = requests.get(RESULT_URL.format(job_id=job_id, img_id=i)) + print("Image", i) + if response.status_code == 200: + # Assuming the endpoint sends the file directly and there's no redirect + p = save_path.format(img_id=i) + with open(p, 'wb') as f: + f.write(response.content) + print(f"Image successfully downloaded and saved to {p}") + else: + # Handle potential error (job not completed, result not available, etc.) + print("Status code: ", response.status_code) + print("Error: ", response.json().get("detail")) # Example usage if __name__ == '__main__': @@ -47,19 +55,35 @@ def get_result(job_id, save_path): parser.add_argument('mask_path', type=str, help='Path to the boolean mask file') # parser.add_argument('coordinates', type=str, help='JSON string of 3D coordinates and labels') parser.add_argument('--prompt', type=str, default='', help='Optional prompt text') - + parser.add_argument('--job_id', type=str, help='Download resources only from a previous job') + parser.add_argument('--resolution', type=int, default=1024, help='Image resolution (W/H)') + parser.add_argument('--num_images', type=int, default=1, help='Number of images to generate') + parser.add_argument('--seed', type=int, default=12345, help='Random seed to use') + parser.add_argument('--num_steps', type=int, default=20, help='Number of processing steps') # Parse the command-line arguments args = parser.parse_args() - - res = submit_image_processing( - args.image_path, - args.mask_path, - # args.coordinates, - args.prompt - ) - print("submitted with response", res) - job_id = res.get('job_id') - print(f"Job submitted. ID: {job_id}") + + if args.job_id: + job_id = args.job_id + else: + res = submit_image_processing( + image_path=args.image_path, + mask_path=args.mask_path, + # args.coordinates, + prompt=args.prompt, + resolution=args.resolution, + num_images=args.num_images, + seed=args.seed, + num_steps=args.num_steps, + ) + j = res.json() + print("submitted with response:", json.dumps(j, indent=4, sort_keys=True)) + if j.get("detail"): + print("Error occurred") + print(j) + sys.exit(1) + job_id = j.get('job_id') + print(f"Job submitted. ID: {job_id}") if job_id: # Wait and check the status until the job is completed @@ -69,7 +93,7 @@ def get_result(job_id, save_path): print(f"Job status: {status_response}") if job_status == 'SUCCESS': # Retrieve and save the result - save_path = 'output_image.png' + save_path = 'output_image_{img_id}.png' get_result(job_id, save_path) break elif job_status == 'FAILURE': diff --git a/gradio_inpaint.py b/gradio_inpaint.py index aeabd6f4..049916cf 100644 --- a/gradio_inpaint.py +++ b/gradio_inpaint.py @@ -33,6 +33,8 @@ def process(input_image_and_mask, prompt, a_prompt, n_prompt, num_samples, image mask_pixel = cv2.resize(input_mask[:, :, 0], (W, H), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255.0 mask_pixel = cv2.GaussianBlur(mask_pixel, (0, 0), mask_blur) + with open("test.txt", "w") as w: + print(mask_pixel, file=w) mask_latent = cv2.resize(mask_pixel, (W // 8, H // 8), interpolation=cv2.INTER_AREA) detected_map = img_raw.copy() diff --git a/server.py b/server.py index e1328671..09c14787 100644 --- a/server.py +++ b/server.py @@ -8,18 +8,24 @@ from celery_app import celery_app, handle_image_processing from celery.result import AsyncResult +from fastapi import HTTPException, status +from fastapi.responses import FileResponse + app = FastAPI() @app.post("/controlnet/") -async def controlnet(background_tasks: BackgroundTasks, - image: UploadFile, # = File(...), - mask: UploadFile, # = File(...), - # coordinates: List[Coordinate], - # prompt: str - ): - prompt = "red drawer" - print("Using prompt: ", prompt) +async def controlnet( + background_tasks: BackgroundTasks, + image: UploadFile, # = File(...), + mask: UploadFile, # = File(...), + # coordinates: List[Coordinate], + prompt: str, + seed: int, + num_images: int, + resolution: int, + num_steps: int +): # Save files temporarily and handle them image_filename = f"/tmp/temp_{uuid4()}_request_img.png" mask_filename = f"/tmp/temp_{uuid4()}_request_mask.png" @@ -31,12 +37,11 @@ async def controlnet(background_tasks: BackgroundTasks, # Create a job ID and start the background task - task = handle_image_processing.delay(image_filename, mask_filename, prompt) + task = handle_image_processing.delay(image_filename, mask_filename, prompt, seed=seed, num_images=num_images, resolution=resolution, num_steps=num_steps) print(f"created job {task.id}") return {"job_id": task.id} - @app.get("/jobs/status/{job_id}") async def get_job_status(job_id: str): result = AsyncResult(job_id, app=celery_app) @@ -47,26 +52,36 @@ async def get_job_status(job_id: str): print("Exception: ", e) return {"job_id": job_id, "status": result.status, "error": str(e)} res_state = {"job_id": job_id, "status": result.status} - if result.status == "PROGRESS": + if result.status == "PROGRESS" or result.status == "SUCCESS": if result.info: print(result.info) res_state.update(result.info) + return res_state -@app.get("/jobs/result/{job_id}") -async def get_result(job_id: str): + +@app.get("/jobs/result/{job_id}/{image_id}") +async def get_result(job_id: str, image_id: int): result = AsyncResult(job_id, app=celery_app) + + if result.state == 'PENDING': + # The job did not start yet + raise HTTPException(status_code=status.HTTP_202_ACCEPTED, detail="Task pending, try again later.") + + elif result.state != 'SUCCESS': + # Something went wrong in the processing + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Task failed with status: {result.state}") + try: res = result.get() + num_images = res["num_images"] # len(res["image_filenames"]) + assert 0 <= image_id < num_images, f"image_id must be between 0 and {num_images}" + image_path = res["image_filenames"][image_id] + return FileResponse(image_path, media_type='image/png', filename=image_path.split("/")[-1]) + except Exception as e: - print("Exception: ", e) - return {"job_id": job_id, "status": result.status, "error": str(e)} - - if result.status != "SUCCESS": - return {"job_id": job_id, "status": result.status, "error": "Result not available or job not completed."} - - image_path = res["image_filename"] - return FileResponse(image_path, media_type='image/png', filename=image_path.split("/")[-1]) + # Handle specific exceptions here if necessary + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) if __name__ == "__main__": From e7606519053e923b2d34318390212955eb9d05fb Mon Sep 17 00:00:00 2001 From: Clemens Marschner Date: Sun, 12 Nov 2023 22:16:35 +0000 Subject: [PATCH 3/4] add ssl --- Makefile | 4 ++++ server.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index a3f48105..298639c2 100644 --- a/Makefile +++ b/Makefile @@ -7,3 +7,7 @@ server: test: python client.py img.png mask.png --prompt "a boatsteg made of glass overlooking a large lake" --seed 1 --num_images 1 --resolution 512 + +cert: + openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/CN=c-marschner.de" + openssl x509 -in cert.pem -outform der -out cert.cer diff --git a/server.py b/server.py index 09c14787..0cab7136 100644 --- a/server.py +++ b/server.py @@ -86,6 +86,6 @@ async def get_result(job_id: str, image_id: int): if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8889) + uvicorn.run(app, host="0.0.0.0", port=8889, ssl_keyfile="./key.pem", ssl_certfile="./cert.pem") From 1d34dc1c03d1fb501bc710a3e94a33c88d8fd66a Mon Sep 17 00:00:00 2001 From: Clemens Marschner Date: Sun, 12 Nov 2023 22:36:24 +0000 Subject: [PATCH 4/4] fix cert gen --- Makefile | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 298639c2..77b7c4e5 100644 --- a/Makefile +++ b/Makefile @@ -8,6 +8,11 @@ server: test: python client.py img.png mask.png --prompt "a boatsteg made of glass overlooking a large lake" --seed 1 --num_images 1 --resolution 512 -cert: - openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/CN=c-marschner.de" +.PHONY: cert +cert.pem key.pem cert.cer: + # openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/CN=c-marschner.de" + openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -config openssl.cnf + openssl x509 -in cert.pem -text -noout openssl x509 -in cert.pem -outform der -out cert.cer + +cert: cert.pem