app.py

import asyncio
import io
from io import BytesIO
import gradio as gr
import requests
import torch
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
from PIL import Image, ImageOps, ImageDraw, ImageFont
import PIL
import replicate
import os

# cuda cpu
device_name = 'cpu'
device = torch.device(device_name)

processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model_clip = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(device)

os.environ['REPLICATE_API_TOKEN'] = '77d68d98d66117c680760543799d97177d6aa722'
model = replicate.models.get("stability-ai/stable-diffusion-inpainting")
version = model.versions.get("c28b92a7ecd66eee4aefcd8a94eb9e7f6c3805d5f06038165407fb5cb355ba67")

sf_prompt_1 = "sunflowers, old bridge, mountain, grass"
sf_neg_prompt_1 = "animal"

sf_prompt_2 = "fire, landscape"
sf_neg_prompt_2 = "animal"

template1 = Image.open("templates/template1.png").resize((512, 512))
template2 = Image.open("templates/template2.png").resize((512, 512))

fontMain = ImageFont.truetype(font="fonts/arial.ttf", size=32)
fontSecond = ImageFont.truetype(font="fonts/arial.ttf", size=18)


def numpy_to_pil(images):
    if images.ndim == 3:
        images = images[None, ...]
    images = (images * 255).round().astype("uint8")

    if images.shape[-1] == 1:
        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
    else:
        pil_images = [Image.fromarray(image) for image in images]

    return pil_images


def get_mask(text, image):
    inputs = processor(
        text=[text], images=[image], padding="max_length", return_tensors="pt"
    ).to(device)

    outputs = model_clip(**inputs)
    mask = torch.sigmoid(outputs.logits).cpu().detach().unsqueeze(-1).numpy()

    mask_pil = numpy_to_pil(mask)[0].resize(image.size)
    return mask_pil


def image_to_byte_array(image: Image) -> bytes:
  imgByteArr = io.BytesIO()
  image.save(imgByteArr, format='PNG')
  #imgByteArr = imgByteArr.getvalue()
  return imgByteArr


def add_template(image, template):
    image.paste(template, (0, 0), mask=template)
    return image


async def predict(prompt, negative_prompt, image, mask_img):
    image = image.convert("RGB").resize((512, 512))
    mask_image = mask_img.convert("RGB").resize((512, 512))
    mask_image = ImageOps.invert(mask_image)

    inputs = {
        # Input prompt
        'prompt': prompt,

        # Specify things to not see in the output
         'negative_prompt': negative_prompt,

        # Inital image to generate variations of. Supproting images size with
        # 512x512
        'image': image_to_byte_array(image),

        # Black and white image to use as mask for inpainting over the image
        # provided. White pixels are inpainted and black pixels are preserved
        'mask': image_to_byte_array(mask_image),

        # Number of images to output. Higher number of outputs may OOM.
        # Range: 1 to 8
        'num_outputs': 1,

        # Number of denoising steps
        # Range: 1 to 500
        'num_inference_steps': 25,

        # Scale for classifier-free guidance
        # Range: 1 to 20
        'guidance_scale': 7.5,

        # Random seed. Leave blank to randomize the seed
        # 'seed': ...,
    }

    output = version.predict(**inputs)

    response = requests.get(output[0])
    img_final = Image.open(BytesIO(response.content))
    mask = mask_image.convert('L')
    PIL.Image.composite(img_final, image, mask)
    return (img_final)


async def predicts(sf_prompt_1, sf_neg_prompt_1, sf_prompt_2, sf_neg_prompt_2, image, image_numpy, mask_img, only_test):
    if only_test:
        img1 = Image.fromarray(image_numpy).convert("RGB").resize((512, 512))
        img2 = Image.fromarray(image_numpy).convert("RGB").resize((512, 512))
        return img1, img2

    task1 = asyncio.create_task(predict(sf_prompt_1, sf_neg_prompt_1, image, mask_img))
    await asyncio.sleep(5)
    task2 = asyncio.create_task(predict(sf_prompt_2, sf_neg_prompt_2, image, mask_img))

    await task1
    await task2

    img1 = task1.result()
    img2 = task2.result()

    return img1, img2


def draw_text(img, template_coords, main_text, second_text):
    x1 = template_coords['x1']
    y1 = template_coords['y1']
    x2 = template_coords['x2']
    y2 = template_coords['y2']

    if '\\n' in main_text:
        main_text = main_text.replace('\\n', '\n')

    if '\\n' in second_text:
        second_text = second_text.replace('\\n', '\n')

    draw = ImageDraw.Draw(img)
    draw.text((x1, y1), main_text, fill=(255, 255, 255), font=fontMain)
    draw.text((x2, y2), second_text, fill=(255, 255, 255), font=fontSecond)


def inference(obj2mask, image_numpy, main_text, second_text, only_test):
    generator = torch.Generator()
    generator.manual_seed(int(52362))

    image = Image.fromarray(image_numpy).convert("RGB").resize((512, 512))

    mask_img = get_mask(obj2mask, image)

    img1, img2 = asyncio.run(predicts(sf_prompt_1, sf_neg_prompt_1, sf_prompt_2, sf_neg_prompt_2, image, image_numpy, mask_img, only_test))

    img1_1 = add_template(img1.copy(), template1.copy())
    img1_2 = add_template(img1.copy(), template2.copy())

    img2_1 = add_template(img2.copy(), template1.copy())
    img2_2 = add_template(img2.copy(), template2.copy())

    template1_coords = {
        'x1': 700/2,
        'y1': 630/2,
        'x2': 420/2,
        'y2': 800/2
    }

    template2_coords = {
        'x1': 30/2,
        'y1': 30/2,
        'x2': 300/2,
        'y2': 740/2
    }

    draw_text(img1_1, template1_coords, main_text, second_text)
    draw_text(img1_2, template2_coords, main_text, second_text)
    draw_text(img2_1, template1_coords, main_text, second_text)
    draw_text(img2_2, template2_coords, main_text, second_text)

    return [img1_1, img1_2, img2_1, img2_2]


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            txt_1 = gr.Textbox(label="Texto principal da propaganda", value="Promoção\nImperdível", lines=2)
            txt_2 = gr.Textbox(label="Texto secundário da propaganda", value="Até 50% para alguns produtos\nEntre em contato com um dos\nnossos vendedores", lines=3)

            mask = gr.Textbox(label="Descrição da imagem", value="shoe")
            intput_img = gr.Image()
            run = gr.Button(value="Gerar")
            chk_test = gr.Checkbox(label='Gerar Prévia')
    with gr.Row():
        with gr.Column():
            output_img1_1 = gr.Image()

        with gr.Column():
            output_img1_2 = gr.Image()

    with gr.Row():
        with gr.Column():
            output_img2_1 = gr.Image()

        with gr.Column():
            output_img2_2 = gr.Image()

    run.click(
        inference,
        inputs=[mask, intput_img, txt_1, txt_2, chk_test],
        outputs=[output_img1_1, output_img1_2, output_img2_1, output_img2_2],
    )

demo.queue(concurrency_count=1)
demo.launch()