added server

Makefile

@@ -0,0 +1,18 @@
+
+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 boatsteg made of glass overlooking a large lake" --seed 1 --num_images 1 --resolution 512
+
+.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

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

backend.py

@@ -0,0 +1,111 @@
+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
+
+        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)
+
+        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 + 1})
+
+        samples, intermediates = ddim_sampler.sample(
+            ddim_steps, num_samples,
+            shape, cond, verbose=True, 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

celery_app.py

@@ -0,0 +1,67 @@
+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, 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):
+             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=num_images,
+            image_resolution=resolution,
+            ddim_steps=num_steps,
+            guess_mode=False,
+            strength=1,
+            scale=9,
+            seed=seed,
+            eta=1,
+            mask_blur=5,
+            update_state_fn=update_state,
+        )
+        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
+        # print(e)
+    finally:
+        # TODO clean-up
+        pass

client.py

@@ -0,0 +1,102 @@
+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}/{img_id}'
+
+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):
+    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__':
+    # 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')
+    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()
+
+    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
+        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_{img_id}.png'
+                get_result(job_id, save_path)
+                break
+            elif job_status == 'FAILURE':
+                print(f"Error: {status_response.get('error')}")
+                break
+            time.sleep(2)

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

gradio_canny.py

@@ -28,24 +28,26 @@
 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':
             detected_map = input_image.copy()
         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,18 +68,18 @@ 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,
                                                      unconditional_conditioning=un_cond)
 
         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)

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