File size: 19,747 Bytes

import cv2
import einops
import numpy as np
import torch
import random
from pytorch_lightning import seed_everything
from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler
from cldm.hack import disable_verbosity, enable_sliced_attention
from datasets.data_utils import * 
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)
import albumentations as A
from omegaconf import OmegaConf
from PIL import Image

save_memory = True
disable_verbosity()
if save_memory:
    enable_sliced_attention()

config = OmegaConf.load('./configs/inference.yaml')
current_model_ckpt = config.pretrained_model
model_config = config.config_file

model = create_model(model_config).cpu()
model.load_state_dict(load_state_dict(current_model_ckpt, location='cuda'))
model = model.cuda()
ddim_sampler = DDIMSampler(model)

def load_model(new_model_ckpt):
    global model, ddim_sampler, current_model_ckpt
    if new_model_ckpt != current_model_ckpt:
        print(f"Loading new model: {new_model_ckpt}")
        model.load_state_dict(load_state_dict(f'/workspace/train-wefadoor-master/anydoor/lightning_logs/version_1/checkpoints/epoch=1-step=2499.ckpt', location='cuda'))
        # model.load_state_dict(load_state_dict(f'/workspace/300k_wefa_boys_slim/lightning_logs/version_0/checkpoints/{new_model_ckpt}', location='cuda'))
        current_model_ckpt = new_model_ckpt
        print("New model loaded successfully.")
    else:
        print("Same model is already loaded, skipping reload.")

def aug_data_mask(image, mask):
    transform = A.Compose([
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.5),
        ])
    transformed = transform(image=image.astype(np.uint8), mask = mask)
    transformed_image = transformed["image"]
    transformed_mask = transformed["mask"]
    return transformed_image, transformed_mask


def process_pairs(ref_image, ref_mask, tar_image, tar_mask):
    # ========= Reference ===========
    # ref expand 
    ref_box_yyxx = get_bbox_from_mask(ref_mask)

    # ref filter mask 
    ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
    masked_ref_image = ref_image * ref_mask_3 + np.ones_like(ref_image) * 255 * (1-ref_mask_3)

    y1,y2,x1,x2 = ref_box_yyxx
    masked_ref_image = masked_ref_image[y1:y2,x1:x2,:]
    ref_mask = ref_mask[y1:y2,x1:x2]


    ratio = np.random.randint(12, 13) / 10
    masked_ref_image, ref_mask = expand_image_mask(masked_ref_image, ref_mask, ratio=ratio)
    ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)

    # to square and resize
    masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False)
    masked_ref_image = cv2.resize(masked_ref_image, (224,224) ).astype(np.uint8)

    ref_mask_3 = pad_to_square(ref_mask_3 * 255, pad_value = 0, random = False)
    ref_mask_3 = cv2.resize(ref_mask_3, (224,224) ).astype(np.uint8)
    ref_mask = ref_mask_3[:,:,0]

    # ref aug 
    masked_ref_image_aug = masked_ref_image #aug_data(masked_ref_image) 

    # collage aug 
    masked_ref_image_compose, ref_mask_compose = masked_ref_image, ref_mask #aug_data_mask(masked_ref_image, ref_mask) 
    masked_ref_image_aug = masked_ref_image_compose.copy()
    ref_mask_3 = np.stack([ref_mask_compose,ref_mask_compose,ref_mask_compose],-1)
    ref_image_collage = sobel(masked_ref_image_compose, ref_mask_compose/255)

    # ========= Target ===========
    tar_box_yyxx = get_bbox_from_mask(tar_mask)
    tar_box_yyxx = expand_bbox(tar_mask, tar_box_yyxx, ratio=[1.1,1.2])

    # crop
    tar_box_yyxx_crop =  expand_bbox(tar_image, tar_box_yyxx, ratio=[1.5, 3])    #1.2 1.6
    tar_box_yyxx_crop = box2squre(tar_image, tar_box_yyxx_crop) # crop box
    y1,y2,x1,x2 = tar_box_yyxx_crop

    cropped_target_image = tar_image[y1:y2,x1:x2,:]
    tar_box_yyxx = box_in_box(tar_box_yyxx, tar_box_yyxx_crop)
    y1,y2,x1,x2 = tar_box_yyxx

    # collage
    ref_image_collage = cv2.resize(ref_image_collage, (x2-x1, y2-y1))
    ref_mask_compose = cv2.resize(ref_mask_compose.astype(np.uint8), (x2-x1, y2-y1))
    ref_mask_compose = (ref_mask_compose > 128).astype(np.uint8)

    collage = cropped_target_image.copy() 
    collage[y1:y2,x1:x2,:] = ref_image_collage

    collage_mask = cropped_target_image.copy() * 0.0
    collage_mask[y1:y2,x1:x2,:] = 1.0

    # the size before pad
    H1, W1 = collage.shape[0], collage.shape[1]
    cropped_target_image = pad_to_square(cropped_target_image, pad_value = 0, random = False).astype(np.uint8)
    collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
    collage_mask = pad_to_square(collage_mask, pad_value = -1, random = False).astype(np.uint8)

    # the size after pad
    H2, W2 = collage.shape[0], collage.shape[1]
    cropped_target_image = cv2.resize(cropped_target_image, (512,512)).astype(np.float32)
    collage = cv2.resize(collage, (512,512)).astype(np.float32)
    collage_mask  = (cv2.resize(collage_mask, (512,512)).astype(np.float32) > 0.5).astype(np.float32)

    masked_ref_image_aug = masked_ref_image_aug  / 255 
    cropped_target_image = cropped_target_image / 127.5 - 1.0
    collage = collage / 127.5 - 1.0 
    collage = np.concatenate([collage, collage_mask[:,:,:1]  ] , -1)

    item = dict(ref=masked_ref_image_aug.copy(), jpg=cropped_target_image.copy(), hint=collage.copy(), extra_sizes=np.array([H1, W1, H2, W2]), tar_box_yyxx_crop=np.array( tar_box_yyxx_crop ) ) 
    return item


def crop_back( pred, tar_image,  extra_sizes, tar_box_yyxx_crop):
    H1, W1, H2, W2 = extra_sizes
    y1,y2,x1,x2 = tar_box_yyxx_crop    
    pred = cv2.resize(pred, (W2, H2))
    m = 5 # maigin_pixel

    if W1 == H1:
        tar_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
        return tar_image

    if W1 < W2:
        pad1 = int((W2 - W1) / 2)
        pad2 = W2 - W1 - pad1
        pred = pred[:,pad1: -pad2, :]
    else:
        pad1 = int((H2 - H1) / 2)
        pad2 = H2 - H1 - pad1
        pred = pred[pad1: -pad2, :, :]

    gen_image = tar_image.copy()
    gen_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
    return gen_image


def inference_single_image(ref_image, ref_mask, tar_image, tar_mask, guidance_scale, seed, steps):
    item = process_pairs(ref_image, ref_mask, tar_image, tar_mask)
    ref = item['ref'] * 255
    tar = item['jpg'] * 127.5 + 127.5
    hint = item['hint'] * 127.5 + 127.5

    hint_image = hint[:,:,:-1]
    hint_mask = item['hint'][:,:,-1] * 255
    hint_mask = np.stack([hint_mask,hint_mask,hint_mask],-1)
    ref = cv2.resize(ref.astype(np.uint8), (512,512))

    seed = random.randint(0, 65535)
    if save_memory:
        model.low_vram_shift(is_diffusing=False)

    ref = item['ref']
    tar = item['jpg'] 
    hint = item['hint']
    num_samples = 1

    control = torch.from_numpy(hint.copy()).float().cuda() 
    control = torch.stack([control for _ in range(num_samples)], dim=0)
    control = einops.rearrange(control, 'b h w c -> b c h w').clone()


    clip_input = torch.from_numpy(ref.copy()).float().cuda() 
    clip_input = torch.stack([clip_input for _ in range(num_samples)], dim=0)
    clip_input = einops.rearrange(clip_input, 'b h w c -> b c h w').clone()

    guess_mode = False
    H,W = 512,512

    cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning( clip_input )]}
    un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([torch.zeros((1,3,224,224))] * num_samples)]}
    shape = (4, H // 8, W // 8)

    if save_memory:
        model.low_vram_shift(is_diffusing=True)

    # ====
    num_samples = 1 #gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
    image_resolution = 512  #gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
    strength = 1  #gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
    guess_mode = False #gr.Checkbox(label='Guess Mode', value=False)
    #detect_resolution = 512  #gr.Slider(label="Segmentation Resolution", minimum=128, maximum=1024, value=512, step=1)
    ddim_steps = steps #gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
    scale = guidance_scale  #gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
    seed = seed  #gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
    eta = 0.0 #gr.Number(label="eta (DDIM)", value=0.0)

    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
    samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
                                                    shape, cond, verbose=False, eta=eta,
                                                    unconditional_guidance_scale=scale,
                                                    unconditional_conditioning=un_cond)
    if 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()#.clip(0, 255).astype(np.uint8)

    result = x_samples[0][:,:,::-1]
    result = np.clip(result,0,255)

    pred = x_samples[0]
    pred = np.clip(pred,0,255)[1:,:,:]
    sizes = item['extra_sizes']
    tar_box_yyxx_crop = item['tar_box_yyxx_crop'] 
    gen_image = crop_back(pred, tar_image, sizes, tar_box_yyxx_crop) 
    return gen_image


import cv2
import numpy as np
import base64
import os
from http.server import BaseHTTPRequestHandler, HTTPServer
import json
from io import BytesIO
from PIL import Image

def base64_to_cv2_image(base64_str):
    img_str = base64.b64decode(base64_str)
    np_img = np.frombuffer(img_str, dtype=np.uint8)
    img = cv2.imdecode(np_img, cv2.IMREAD_COLOR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img

def base64_to_pil_image(base64_str):
    img_data = base64.b64decode(base64_str)
    img = Image.open(BytesIO(img_data))
    return img

def pil_image_to_np_array(pil_img, target_index):
    np_array = np.array(pil_img)
    return (np_array == target_index).astype(np.uint8)

def image_to_base64(img):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    _, buffer = cv2.imencode('.jpg', img)
    base64_str = base64.b64encode(buffer).decode("utf-8")
    return base64_str

class RequestHandler(BaseHTTPRequestHandler):
    API_KEY = "xiCQTaoQKXUNATzuFLWRgtoJKiFXiDGvnk"

    def _set_response(self, status_code=200, content_type='application/json'):
        self.send_response(status_code)
        self.send_header('Content-type', content_type)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'X-API-Key, Content-Type')
        self.end_headers()

    def do_OPTIONS(self):
        self._set_response(204)

    def do_GET(self):
        self._set_response(405)
        self.wfile.write(b'{"error": "GET method not allowed."}')

    def handle_not_supported_method(self):
        self._set_response(405)
        self.wfile.write(b'{"error": "Method not supported."}')

    def do_PUT(self):
        self.handle_not_supported_method()

    def do_DELETE(self):
        self.handle_not_supported_method()

    def do_PATCH(self):
        self.handle_not_supported_method()
        
    def do_POST(self):
        print("Received POST request...")
        received_api_key = self.headers.get('X-API-Key')
        
        if received_api_key != self.API_KEY:
            self._set_response(401)
            self.wfile.write(b'{"error": "Invalid API key"}')
            print("Invalid API key")
            return

        content_length = int(self.headers['Content-Length'])
        print(f"Content Length: {content_length}")
        
        if content_length:
            post_data = self.rfile.read(content_length)
            print("Data received")
            try:
                data = json.loads(post_data.decode('utf-8'))
                print("Processing data")
                
                model_name = data.get('model', 'default_model.ckpt')
                model_ckpt_map = {
                    'boys': 'boys.ckpt',
                    'men': 'men.ckpt',
                    'women': 'women.ckpt',
                    'girls': 'girls.ckpt'
                }
                new_model_ckpt = model_ckpt_map.get(model_name, current_model_ckpt)
                load_model(new_model_ckpt)

                seed = int(data.get('seed'))
                steps = int(data.get('steps'))
                guidance_scale = float(data.get('guidance_scale'))

                ref_image = base64_to_cv2_image(data['ref_image'])
                tar_image = base64_to_cv2_image(data['tar_image'])

                ref_mask_img = base64_to_cv2_image(data['ref_mask'])
                ref_mask = cv2.cvtColor(ref_mask_img, cv2.COLOR_RGB2GRAY)
                ref_mask = (ref_mask > 128).astype(np.uint8)

                tar_mask_img = base64_to_cv2_image(data['tar_mask'])
                tar_mask = cv2.cvtColor(tar_mask_img, cv2.COLOR_RGB2GRAY)
                tar_mask = (tar_mask > 128).astype(np.uint8)

                gen_image = inference_single_image(ref_image, ref_mask, tar_image, tar_mask, guidance_scale, seed, steps)
                gen_image_base64 = image_to_base64(gen_image)

                self.send_response(200)
                self.send_header('Content-Type', 'image/jpeg')
                self.end_headers()
                self.wfile.write(base64.b64decode(gen_image_base64))

                print("Sent image response")

            except Exception as e:
                print(f"An error occurred: {e}")
                self._set_response(500)
                error_data = json.dumps({'error': str(e)}).encode('utf-8')
                self.wfile.write(error_data)
                print("Sent error response")

        else:
            print("No data received in POST request.")
            self._set_response(400)
            error_data = json.dumps({'error': 'No data received'}).encode('utf-8')
            self.wfile.write(error_data)
            print("Sent error response")

def run(server_class=HTTPServer, handler_class=RequestHandler, port=8084):
    server_address = ('', port)
    httpd = server_class(server_address, handler_class)
    print(f"Starting HTTP server on port {port}")
    httpd.serve_forever()

if __name__ == "__main__":
    run()

# class RequestHandler(BaseHTTPRequestHandler):
#     API_KEY = "xiCQTaoQKXUNATzuFLWRgtoJKiFXiDGvnk"

#     def _set_response(self, status_code=200, content_type='application/json'):
#         self.send_response(status_code)
#         self.send_header('Content-type', content_type)
#         self.send_header('Access-Control-Allow-Origin', '*')
#         self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
#         self.send_header('Access-Control-Allow-Headers', 'X-API-Key, Content-Type')
#         self.end_headers()

#     def do_OPTIONS(self):
#         self._set_response(204)  # No content to send back for OPTIONS request

#     def do_GET(self):
#         # If needed, define handling for GET or send a 405 if it's not supported
#         self._set_response(405)
#         self.wfile.write(b'{"error": "GET method not allowed."}')

#     def handle_not_supported_method(self):
#         self._set_response(405)
#         self.wfile.write(b'{"error": "Method not supported."}')

#     def do_PUT(self):
#         self.handle_not_supported_method()

#     def do_DELETE(self):
#         self.handle_not_supported_method()

#     def do_PATCH(self):
#         self.handle_not_supported_method()
        
#     def do_POST(self):
#         print("Received POST request...")
#         received_api_key = self.headers.get('X-API-Key')
#         # Check if the API key is correct
#         if received_api_key != self.API_KEY:
#             # If the API key is incorrect, respond with 401 Unauthorized
#             self._set_response(401)
#             self.wfile.write(b'{"error": "Invalid API key"}')
#             print("Invalid API key")
#             return

#         content_length = int(self.headers['Content-Length'])
#         print(f"Content Length: {content_length}")
        
#         if content_length:
#             post_data = self.rfile.read(content_length)
#             print("Data received")
#             try:
#                 data = json.loads(post_data.decode('utf-8'))
#                 print("Processing data")
#                 # print(data)

#                 seed = int(data.get('seed'))
#                 steps = int(data.get('steps'))
#                 guidance_scale = float(data.get('guidance_scale'))

#                 ref_image = base64_to_cv2_image(data['ref_image'])
#                 tar_image = base64_to_cv2_image(data['tar_image'])
#                 # print(seed)
#                 # print(steps)
#                 # print(guidance_scale) 
#                 # Process reference mask
#                 ref_mask_img = base64_to_cv2_image(data['ref_mask'])
#                 ref_mask = cv2.cvtColor(ref_mask_img, cv2.COLOR_RGB2GRAY)
#                 ref_mask = (ref_mask > 128).astype(np.uint8)

#                 # Process target mask
#                 tar_mask_img = base64_to_cv2_image(data['tar_mask'])
#                 tar_mask = cv2.cvtColor(tar_mask_img, cv2.COLOR_RGB2GRAY)
#                 tar_mask = (tar_mask > 128).astype(np.uint8)

#                 output_dir = '/work/ADOOR_ACE/test_out'
#                 os.makedirs(output_dir, exist_ok=True)

#                 # Save reference and target images
#                 cv2.imwrite(os.path.join(output_dir, 'out_ref_image.jpg'), cv2.cvtColor(ref_image, cv2.COLOR_RGB2BGR))
#                 cv2.imwrite(os.path.join(output_dir, 'out_tar_image.jpg'), cv2.cvtColor(tar_image, cv2.COLOR_RGB2BGR))

#                 # Save reference mask
#                 ref_mask_img_to_save = (ref_mask * 255).astype(np.uint8)
#                 cv2.imwrite(os.path.join(output_dir, 'out_ref_mask.jpg'), ref_mask_img_to_save)

#                 # Save target mask
#                 tar_mask_img_to_save = (tar_mask * 255).astype(np.uint8)
#                 cv2.imwrite(os.path.join(output_dir,'out_tar_mask.jpg'), tar_mask_img_to_save)

#                 gen_image = inference_single_image(ref_image, ref_mask, tar_image, tar_mask, guidance_scale, seed, steps)
#                 gen_image_base64 = image_to_base64(gen_image)

#                 self.send_response(200)
#                 self.send_header('Content-Type', 'image/jpeg')
#                 self.end_headers()
#                 self.wfile.write(base64.b64decode(gen_image_base64))

#                 print("Sent image response")

#             except Exception as e:
#                 print(f"An error occurred: {e}")
#                 self._set_response(500)
#                 error_data = json.dumps({'error': str(e)}).encode('utf-8')
#                 self.wfile.write(error_data)
#                 print("Sent error response")

#         else:
#             print("No data received in POST request.")
#             self._set_response(400)
#             error_data = json.dumps({'error': 'No data received'}).encode('utf-8')
#             self.wfile.write(error_data)
#             print("Sent error response")

        

# def run(server_class=HTTPServer, handler_class=RequestHandler, port=8084):
#     server_address = ('', port)
#     httpd = server_class(server_address, handler_class)
#     print(f"Starting HTTP server on port {port}")
#     httpd.serve_forever()

# if __name__ == "__main__":
#     run()