model/dongba_dreamer.py

import os
import cv2
import torch
import numpy as np
from omegaconf import OmegaConf
from random import randint
import json
# import codecs
from . import dongba_sd_helper as sd_helper
from PIL import Image
from .dongda_gpt_helper import GPTHelper
from .omegaconf_utils import load_from_config
import gradio as gr

def _overlay_composition(img, x, y, w, h, alpha, clr, composition_image):
    # img to weights
    paint = 1.0 - (img.astype(float) / 255.0).mean(axis=2)

    # clip
    mask = (paint > 0.25)
    paint[mask] = 1.0
    paint[~mask] = 0.0

    # # gamma before resize
    # paint = np.power(paint, 1.0 / 2.2)

    t = int(min(max(composition_image.shape[0] * y, 0), composition_image.shape[0] - 2))
    b = int(min(max(composition_image.shape[0] * (y + h), 0), composition_image.shape[0] - 1))
    l = int(min(max(composition_image.shape[1] * x, 0), composition_image.shape[1] - 2))
    r = int(min(max(composition_image.shape[1] * (x + w), 0), composition_image.shape[1] - 1))

    to_w = int(r - l)
    to_h = int(b - t)
    paint = cv2.resize(paint, (to_w, to_h))[:, :, None]

    # # gamma after resize
    # paint = np.power(paint, 2.2)

    # blending
    block = composition_image[t:b, l:r, :]
    block = block * (1.0 - paint * alpha) + (clr[None, None, :] * paint * alpha)
    composition_image[t:b, l:r, :] = block
    return composition_image
    

class DongbaDreamer():
    def __init__(self, config) -> None:
        self.setup_config(config)
        self.reset_logs()

    def reset_logs(self):
        self.logs = 'Dongba Dreamer\n'

    def setup_config(self, config):
        self.config = config
        config.gpt.yaml_dir = config.yaml_dir

        # gpt
        self.gpt_helper = GPTHelper(config.gpt)
        self.gpt_helper.log = self.log

        # sd
        if config.get('generator_sd', None):
            self.generator_sd = sd_helper.SDHelper(config.generator_sd)

        if config.get('controlnet_sd', None):
            self.controlnet_sd = sd_helper.SDHelper(config.controlnet_sd)

    def get_logs(self):
        return  self.logs
    
    def log(self, text, end='\n'):
        print(f'{text}')
        self.logs += text + end

    ##################################################
    def composite_image(self, canvas_width, canvas_height):
        dat_dir = self.config.dat_dir
        composition_image = np.zeros((canvas_height, canvas_width, 3))

        x_max = max([item['x'] for item in self.info['composition']])
        y_max = max([item['y'] for item in self.info['composition']])

        x_max = x_max * 1.2
        y_max = y_max * 1.2

        for item in self.info['composition']:
            idx = item['idx']
            img_dir = os.path.join(dat_dir, f'{idx+1:04d}')
            if not os.path.exists(img_dir):
                self.log(f'{img_dir} not exist!')
                continue

            img_fns = [fn for fn in os.listdir(img_dir) if fn.endswith('.png')]
            if len(img_fns) == 0:
                self.log(f'{img_dir} is empty!')
                continue

            x = item['x'] / x_max
            y = item['y'] / y_max
            w = 180 / canvas_width
            h = 180 / canvas_height
            x = x - w / 2
            y = y - h / 2
            
            select_i = randint(0, len(img_fns)-1)
            img_fn = os.path.join(img_dir, img_fns[select_i])

            img = cv2.imread(img_fn, cv2.IMREAD_UNCHANGED)
            img, alpha = img[:, :, :3], img[:, :, 3:]
            clr = np.array([1.0, 1.0, 1.0])
            img = 255 - img
            img[img < 250] = 0
            alpha = 1.0
            composition_image = _overlay_composition(img, x, y, w, h, alpha, clr, composition_image)

        composition_image = (composition_image.clip(0, 1) * 255).astype(np.uint8)
        return composition_image

    ##################################################
    def process_words(self, image_topic, canvas_width, canvas_height, num_words=0):
        if image_topic is None or len(image_topic) == 0:
            return []

        self.log('----------------------------------------')
        self.log(f'这幅画的主题是：')
        self.log(f'{image_topic}')
        self.log(f'分辨率是 {canvas_width}x{canvas_height}')
        self.log('----------------------------------------')

        self.log(f'<GPT> 生成 stable diffusion prompt: ')
        image_prompt = self.gpt_helper.query_image_prompt(image_topic)
        image_prompt = 'Joan Miro style abstract painting. ' + image_prompt
        self.log(f'{image_prompt}')
        self.log('----------------------------------------')
        
        keywords = self.gpt_helper.query_keywords(image_topic)
        self.log(f'<GPT> 画面关键词有：{keywords}')
        self.log('----------------------------------------')

        ##################################################
        self.log(f'查询东巴文数据库：')
        words_to_use = []
        for keyword in keywords:
            words = self.gpt_helper.query_in_faiss_db(keyword)
            if len(words) > 0:
                select_i = 0
                self.log(f"- {keyword}: {[w['word'] for w in words]}, 选择：{words[select_i]['word']}")
                words_to_use.append(words[select_i])

        # keywords_to_query = ', '.join([w['word'] for w in words_to_use])
        keywords_to_query = ', '.join([f"{w['idx']}:{w['word']}" for w in words_to_use])
        self.log('----------------------------------------')

        ##################################################
        self.info = {
            'image_topic': image_topic,
            'keywords': keywords,
            'words_to_use': words_to_use,
            'keywords_to_query': keywords_to_query,
        }

        # info_fn = os.path.join(self.config.work_dir, 'info.json')
        # with codecs.open(info_fn, 'w', encoding='utf-8') as fp:
        #     json.dump(self.info, fp, indent=4, ensure_ascii=False)

        ##################################################
        self.log('<gpt> 生成构图...', end='')
        composition_txt = self.gpt_helper.query_composition(image_topic, keywords_to_query, canvas_width, canvas_height, num_words=num_words)
        self.log(composition_txt)

        composition_txt = composition_txt.replace('```', '')
        composition_txt = composition_txt.replace('json', '')

        # with open(os.path.join(self.config.work_dir, 'composition.json'), 'w') as fp:
        #     fp.write(composition_txt)

        composition = json.loads(composition_txt.replace('/n', ''))

        self.info.update({
            'composition': composition,
        })
        # with codecs.open(info_fn, 'w', encoding='utf-8') as fp:
        #     json.dump(self.info, fp, indent=4, ensure_ascii=False)

        ##################################################
        word_images = []
        self.log('构图由远及近分别是：')

        for word_i, item in enumerate(composition):
            idx = item['idx']
            img_dir = os.path.join(self.config.dat_dir, f'{idx+1:04d}')
            if not os.path.exists(img_dir):
                self.log(f'{img_dir} not exist!')
                continue

            self.log(f'- {item["name"]}')

            img_fns = [fn for fn in os.listdir(img_dir) if fn.endswith('.png')]

            # select_i = randint(0, len(img_fns)-1)
            select_i = 0
            img_fn = os.path.join(img_dir, img_fns[select_i])
            img = cv2.imread(img_fn)

            self.info['composition'][word_i]['select_i'] = select_i
            word_images.append(Image.fromarray(img))
        self.log('----------------------------------------')

        self.word_images = word_images

        # composite
        composition_image = self.composite_image(canvas_width, canvas_height)

        return {
            'word_images': word_images,
            'image_prompt': image_prompt,
            'composition_image': composition_image,
        }

    ##################################################
    def process_sd(self, image_prompt, composition_image):
        self.log(f'<sd> 生成图片: {image_prompt}')

        if self.config.condition_type == 'mixed_canny':
            reference_image = self.generator_sd.forward(image_prompt)[0]
            composition_image = composition_image
            
            control_canny = sd_helper.make_merged_canny(composition_image, reference_image, 0.5)
            control_image = Image.fromarray(control_canny)
            # control_image.save(os.path.join(work_dir, 'control_canny.jpg'))

            self.log('<sd> 注入灵魂...')
            rlt_image = self.controlnet_sd.forward(image_prompt, control_image=control_image)
            # image.save(os.path.join(work_dir, "moon-v1.5-canny.png"))

            images = [*rlt_image, control_image, 
                      reference_image,
                      Image.fromarray(composition_image)]
        else:
            control_image = Image.fromarray(composition_image)
            rlt_image = self.controlnet_sd.forward(image_prompt, control_image=control_image)
            images = [*rlt_image, control_image]

        return images