app.py

import os
# os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
# os.environ['CUDA_VISIBLE_DEVICES'] = '2'
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "caching_allocator"
import gradio as gr
import numpy as np
from models import make_inpainting
import utils

from transformers import MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
from PIL import Image
import requests
from transformers import pipeline
import torch
import random
import io
import base64
import json
from diffusers import DiffusionPipeline
from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline
from diffusers import StableDiffusionUpscalePipeline
from diffusers import LDMSuperResolutionPipeline
import cv2
import onnxruntime
from split_image import split

def removeFurniture(input_img1,
            input_img2,
            positive_prompt,
            negative_prompt,
            num_of_images,
            resolution
            ):

    print("removeFurniture")
    HEIGHT = resolution
    WIDTH = resolution

    input_img1 = input_img1.resize((resolution, resolution))
    input_img2 = input_img2.resize((resolution, resolution))

    canvas_mask = np.array(input_img2)
    mask = utils.get_mask(canvas_mask)

    print(input_img1, mask, positive_prompt, negative_prompt)

    retList=  make_inpainting(positive_prompt=positive_prompt,
                               image=input_img1,
                               mask_image=mask,
                               negative_prompt=negative_prompt,
                               num_of_images=num_of_images,
                               resolution=resolution
                               )
    # add the rest up to 10
    while (len(retList)<10):
        retList.append(None)

    return retList

def imageToString(img):

    output = io.BytesIO()
    img.save(output, format="png")
    return output.getvalue()

def segmentation(img):
    print("segmentation")

    # semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
    pipe = pipeline("image-segmentation", "facebook/maskformer-swin-large-ade")    
    results = pipe(img)
    for p in results:
        p['mask'] = utils.image_to_byte_array(p['mask'])
        p['mask'] = base64.b64encode(p['mask']).decode("utf-8")
    #print(results)
    return json.dumps(results)
    




def upscale1(image, prompt):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("upscale1", device, image, prompt)    
    
    # image.thumbnail((512, 512))
    # print("resize",image)

    torch.backends.cuda.matmul.allow_tf32 = True

    pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler",                                                           
                                                          torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                                                          use_safetensors=True)
    # pipe = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
    pipe = pipe.to(device)
    pipe.enable_attention_slicing()
    pipe.enable_xformers_memory_efficient_attention()
    # pipe.enable_xformers_memory_efficient_attention(attention_op=xformers.ops.MemoryEfficientAttentionFlashAttentionOp)
    # Workaround for not accepting attention shape using VAE for Flash Attention
    pipe.vae.enable_xformers_memory_efficient_attention()

    ret = pipe(prompt=prompt, 
                   image=image,
                   num_inference_steps=10,
                   guidance_scale=0)
    print("ret",ret)
    upscaled_image = ret.images[0]
    print("up",upscaled_image)

    return upscaled_image

def upscale2(image, prompt):
    print("upscale2",image,prompt)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("device",device)

    pipe = LDMSuperResolutionPipeline.from_pretrained("CompVis/ldm-super-resolution-4x-openimages", torch_dtype=torch.float16)
    pipe = pipe.to(device)
    pipe.enable_attention_slicing()
    pipe.enable_xformers_memory_efficient_attention(attention_op=xformers.ops.MemoryEfficientAttentionFlashAttentionOp)
    # Workaround for not accepting attention shape using VAE for Flash Attention
    pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)

    upscaled_image = pipe(image, num_inference_steps=10, eta=1).images[0]
    return upscaled_image

def convert_pil_to_cv2(image):
    # pil_image = image.convert("RGB")
    open_cv_image = np.array(image)
    # RGB to BGR
    open_cv_image = open_cv_image[:, :, ::-1].copy()
    return open_cv_image

def inference(model_path: str, img_array: np.array) -> np.array:
    options = onnxruntime.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    ort_session = onnxruntime.InferenceSession(model_path, options)
    ort_inputs = {ort_session.get_inputs()[0].name: img_array}
    ort_outs = ort_session.run(None, ort_inputs)

    return ort_outs[0]

def post_process(img: np.array) -> np.array:
    # 1, C, H, W -> C, H, W
    img = np.squeeze(img)
    # C, H, W -> H, W, C
    img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
    return img

def pre_process(img: np.array) -> np.array:
    # H, W, C -> C, H, W
    img = np.transpose(img[:, :, 0:3], (2, 0, 1))
    # C, H, W -> 1, C, H, W
    img = np.expand_dims(img, axis=0).astype(np.float32)
    return img

def upscale3(image):
    print("upscale3",image)

    model_path = f"up_models/modelx4.ort"
    img = convert_pil_to_cv2(image)
    
    # if img.ndim == 2:
    #     print("upscale3","img.ndim == 2")
    #     img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    # if img.shape[2] == 4:
    #     print("upscale3","img.shape[2] == 4")
    #     alpha = img[:, :, 3]  # GRAY
    #     alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
    #     alpha_output = post_process(inference(model_path, pre_process(alpha)))  # BGR
    #     alpha_output = cv2.cvtColor(alpha_output, cv2.COLOR_BGR2GRAY)  # GRAY

    #     img = img[:, :, 0:3]  # BGR
    #     image_output = post_process(inference(model_path, pre_process(img)))  # BGR
    #     image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2BGRA)  # BGRA
    #     image_output[:, :, 3] = alpha_output

    # print("upscale3","img.shape[2] == 3")
    image_output = post_process(inference(model_path, pre_process(img)))  # BGR

    return image_output



def split_image(im, rows, cols, should_square, should_quiet=False):
    im_width, im_height = im.size
    row_width = int(im_width / cols)
    row_height = int(im_height / rows)
    name = "image"
    ext = ".png"
    name = os.path.basename(name)
    images = []
    if should_square:
        min_dimension = min(im_width, im_height)
        max_dimension = max(im_width, im_height)
        if not should_quiet:
            print("Resizing image to a square...")
            print("Determining background color...")
        bg_color = split.determine_bg_color(im)
        if not should_quiet:
            print("Background color is... " + str(bg_color))
        im_r = Image.new("RGBA" if ext == "png" else "RGB",
                         (max_dimension, max_dimension), bg_color)
        offset = int((max_dimension - min_dimension) / 2)
        if im_width > im_height:
            im_r.paste(im, (0, offset))
        else:
            im_r.paste(im, (offset, 0))
        im = im_r
        row_width = int(max_dimension / cols)
        row_height = int(max_dimension / rows)
    n = 0
    for i in range(0, rows):
        for j in range(0, cols):
            box = (j * row_width, i * row_height, j * row_width +
                   row_width, i * row_height + row_height)
            outp = im.crop(box)
            outp_path = name + "_" + str(n) + ext
            if not should_quiet:
                print("Exporting image tile: " + outp_path)
            images.append(outp)
            n += 1
    return [img for img in images]

def upscale_image(img, rows, up_factor, cols, seed, prompt, negative_prompt, xformers, cpu_offload, attention_slicing, enable_custom_sliders=False, guidance=7, iterations=50):
    
    if up_factor==2:
        model_id = "stabilityai/sd-x2-latent-upscaler"
        try:
            pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
        except:
            pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16, local_files_only=True)

    if up_factor==4:
        model_id = "stabilityai/stable-diffusion-x4-upscaler"

        try:
            pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
        except:
            pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16, local_files_only=True)            


    pipeline = pipeline.to("cuda")
    if xformers:
        pipeline.enable_xformers_memory_efficient_attention()
    else:
        pipeline.disable_xformers_memory_efficient_attention()
    if cpu_offload:
        try:
            pipeline.enable_sequential_cpu_offload()
        except:
            pass
    if attention_slicing:
        pipeline.enable_attention_slicing()
    else:
        pipeline.disable_attention_slicing()
    img = Image.fromarray(img)
    # load model and scheduler
    if seed==-1:
        generator = torch.manual_seed(random.randint(0, 9999999))
    else:
        generator = torch.manual_seed(seed)
    
    original_width, original_height = img.size
    max_dimension = max(original_width, original_height)
    tiles = split_image(img, rows, cols, True, False)
    ups_tiles = []
    i = 0
    for x in tiles:
        i=i+1
        if enable_custom_sliders:
            ups_tile = pipeline(prompt=prompt,negative_prompt=negative_prompt,guidance_scale=guidance, num_inference_steps=iterations, image=x.convert("RGB"),generator=generator).images[0]
        else:
            ups_tile = pipeline(prompt=prompt,negative_prompt=negative_prompt, image=x.convert("RGB"),generator=generator).images[0]
        ups_tiles.append(ups_tile)
        
    # Determine the size of the merged upscaled image
    total_width = 0
    total_height = 0
    side = 0
    for ups_tile in ups_tiles:
        side = ups_tile.width
        break
    for x in tiles:
        tsize = x.width
        break

    ups_times = abs(side/tsize)
    new_size = (max_dimension * ups_times, max_dimension * ups_times)
    total_width = cols*side
    total_height = rows*side

    # Create a blank image with the calculated size
    merged_image = Image.new("RGB", (total_width, total_height))

    # Paste each upscaled tile into the blank image
    current_width = 0
    current_height = 0
    maximum_width = cols*side
    for ups_tile in ups_tiles:
        merged_image.paste(ups_tile, (current_width, current_height))
        current_width += ups_tile.width
        if current_width>=maximum_width:
            current_width = 0
            current_height = current_height+side

    # Using the center of the image as pivot, crop the image to the original dimension times four
    crop_left = (new_size[0] - original_width * ups_times) // 2
    crop_upper = (new_size[1] - original_height * ups_times) // 2
    crop_right = crop_left + original_width * ups_times
    crop_lower = crop_upper + original_height * ups_times
    final_img = merged_image.crop((crop_left, crop_upper, crop_right, crop_lower))

    # The resulting image should be identical to the original image in proportions / aspect ratio, with no loss of elements.
    # Save the merged image
    return final_img

    
def upscale( image, prompt, negative_prompt, rows, up_factor, guidance, iterations, xformers_input, cpu_offload_input, attention_slicing_input):
    print("upscale", prompt, negative_prompt, rows, up_factor, guidance, iterations, xformers_input, cpu_offload_input, attention_slicing_input)   
    return upscale_image(img=image, 
                         rows=rows,cols=rows,
                         up_factor=up_factor,
                         seed=-1,
                         prompt=prompt,
                         negative_prompt=negative_prompt, 
                         enable_custom_sliders=True,
                         xformers=xformers_input, 
                         cpu_offload=cpu_offload_input, 
                         attention_slicing=attention_slicing_input,
                         guidance=guidance,
                         iterations=iterations)

# modes = {
#     '1': '1',
#     'img2img': 'Image to Image',
#     'inpaint': 'Inpainting',
#     'upscale4x': 'Upscale 4x',
# }



with gr.Blocks() as app:    
    gr.HTML(
        f"""         
            Running on <b>{"GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"}</b>
          </div>
        """
    )

    with gr.Row():

        with gr.Column():
            gr.Button("FurnituRemove").click(removeFurniture, 
                                        inputs=[gr.Image(label="img", type="pil"),
                                                gr.Image(label="mask", type="pil"),
                                                gr.Textbox(label="positive_prompt",value="empty room"),
                                                gr.Textbox(label="negative_prompt",value=""),
                                                gr.Number(label="num_of_images",value=2),
                                                gr.Number(label="resolution",value=512)
                                                ], 
                                        outputs=[
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image(),
                                                gr.Image()])
        
        with gr.Column():  
            gr.Button("Segmentation").click(segmentation, inputs=gr.Image(type="pil"), outputs=gr.JSON())

        with gr.Column():
            gr.Button("Upscale").click(
                upscale, 
                inputs=[
                    gr.Image(label="Source Image to upscale"),
                    gr.Textbox(label="prompt",value="empty room"),
                    gr.Textbox(label="negative prompt",value="jpeg artifacts, lowres, bad quality, watermark, text"),
                    gr.Number(value=2, label="Tile grid dimension amount (number of rows and columns) - X by X "),
                    gr.Slider(2, 4, 2, step=2, label='Upscale 2 or 4'),
                    gr.Slider(2, 15, 7, step=1, label='Guidance Scale: How much the AI influences the Upscaling.'),
                    gr.Slider(2, 100, 10, step=1, label='Number of Iterations'),
                    gr.Checkbox(value=True,label="Enable Xformers memory efficient attention"),                    
                    gr.Checkbox(value=True,label="Enable sequential CPU offload"),
                    gr.Checkbox(value=True,label="Enable attention slicing")
                    ], 
                outputs=gr.Image())


# app.queue()
app.launch()

# UP 1