PhotoMaker と OpenVINO を使用したテキストからの画像生成#

from pathlib import Path 

if not Path("PhotoMaker").exists():
     !git clone https://github.com/TencentARC/PhotoMaker.git

%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu\ transformers "torch>=2.1" "diffusers>=0.26" "gradio>=4.19" "openvino>=2024.0.0" torchvision "peft==0.6.2" "nncf>=2.9.0" "protobuf==3.20.3"

adapter_id = "TencentARC/PhotoMaker" 
base_model_id = "SG161222/RealVisXL_V3.0" 

TEXT_ENCODER_OV_PATH = Path("model/text_encoder.xml") 
TEXT_ENCODER_2_OV_PATH = Path("model/text_encoder_2.xml") 
UNET_OV_PATH = Path("model/unet.xml") 
ID_ENCODER_OV_PATH = Path("model/id_encoder.xml") 
VAE_DECODER_OV_PATH = Path("model/vae_decoder.xml")

import torch 
import numpy as np 
import os 
from PIL import Image 
from pathlib import Path 
from PhotoMaker.photomaker.model import PhotoMakerIDEncoder 
from PhotoMaker.photomaker.pipeline import PhotoMakerStableDiffusionXLPipeline 
from diffusers import EulerDiscreteScheduler 
import gc 

trigger_word = "img" 

def load_original_pytorch_pipeline_components(photomaker_path: str, base_model_id: str):
    # ベースモデルをロード 
    pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(base_model_id, use_safetensors=True).to("cpu") 

    # PhotoMaker チェックポイントをロード 
    pipe.load_photomaker_adapter( 
        os.path.dirname(photomaker_path), 
        subfolder="", 
        weight_name=os.path.basename(photomaker_path), 
        trigger_word=trigger_word, 
    ) 
    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config) 
    pipe.fuse_lora() 
    gc.collect() 
    return pipe

from huggingface_hub import hf_hub_download 

photomaker_path = hf_hub_download(repo_id=adapter_id, filename="photomaker-v1.bin", repo_type="model") 

pipe = load_original_pytorch_pipeline_components(photomaker_path, base_model_id)

import openvino as ov 
import nncf 

def flattenize_inputs(inputs): 
    """ 
    Helper function for resolve nested input structure (e.g. lists or tuples of tensors) 
    """ 
    flatten_inputs = [] 
    for input_data in inputs: 
        if input_data is None: 
            continue 
        if isinstance(input_data, (list, tuple)): 
            flatten_inputs.extend(flattenize_inputs(input_data)) 
        else: 
            flatten_inputs.append(input_data) 
    return flatten_inputs 

dtype_mapping = { 
    torch.float32: ov.Type.f32, 
    torch.float64: ov.Type.f64, 
    torch.int32: ov.Type.i32, 
    torch.int64: ov.Type.i64, 
    torch.bool: ov.Type.boolean, 
} 

def prepare_input_info(input_dict): 
    """ 
    Helper function for preparing input info (shapes and data types) for conversion based on example inputs 
    """ 
    flatten_inputs = flattenize_inputs(input_dict.values()) 
    input_info = [] 
    for input_data in flatten_inputs: 
        updated_shape = list(input_data.shape) 
        if input_data.ndim == 5: 
            updated_shape[1] = -1 
        input_info.append((dtype_mapping[input_data.dtype], updated_shape)) 
    return input_info 

def convert(model: torch.nn.Module, xml_path: str, example_input, input_info): 
    """ 
    Helper function for converting PyTorch model to OpenVINO IR 
    """ 
    xml_path = Path(xml_path) 
    if not xml_path.exists(): 
        xml_path.parent.mkdir(parents=True, exist_ok=True) 
        with torch.no_grad(): 
            ov_model = ov.convert_model(model, example_input=example_input, input=input_info) 
        ov_model = nncf.compress_weights(ov_model) 
        ov.save_model(ov_model, xml_path) 

        del ov_model 
        torch._C._jit_clear_class_registry() 
        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() 
        torch.jit._state._clear_class_state()

from diffusers.utils import load_image 

prompt = "sci-fi, closeup portrait photo of a man img in Iron man suit, face" 
negative_prompt = "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth" 
generator = torch.Generator("cpu").manual_seed(42) 

input_id_images = [] 
original_image = load_image("./PhotoMaker/examples/newton_man/newton_0.jpg") 
input_id_images.append(original_image) 

## パラメーター設定 
num_steps = 20 
style_strength_ratio = 20 
start_merge_step = int(float(style_strength_ratio) / 100 * num_steps) 
if start_merge_step > 30: 
    start_merge_step = 30 

images = pipe( 
    prompt=prompt, 
    input_id_images=input_id_images, 
    negative_prompt=negative_prompt, 
    num_images_per_prompt=1, 
    num_inference_steps=num_steps, 
    start_merge_step=start_merge_step, 
    generator=generator, 
).images

import matplotlib.pyplot as plt 

def visualize_results(orig_img: Image.Image, output_img: Image.Image): 
    """ 
    Helper function for pose estimationresults visualization 

    Parameters: 
        orig_img (Image.Image): original image 
        output_img (Image.Image): processed image with PhotoMaker 
    Returns: 
        fig (matplotlib.pyplot.Figure): matplotlib generated figure 
    """ 
    orig_img = orig_img.resize(output_img.size) 
    orig_title = "Original image" 
    output_title = "Output image" 
    im_w, im_h = orig_img.size 
    is_horizontal = im_h < im_w 
    fig, axs = plt.subplots( 
        2 if is_horizontal else 1, 
        1 if is_horizontal else 2, 
        sharex="all", 
        sharey="all", 
    ) 
    fig.suptitle(f"Prompt: '{prompt}'", fontweight="bold") 
    fig.patch.set_facecolor("white") 
    list_axes = list(axs.flat) 
    for a in list_axes: 
        a.set_xticklabels([]) 
        a.set_yticklabels([]) 
        a.get_xaxis().set_visible(False) 
        a.get_yaxis().set_visible(False) 
        a.grid(False) 
    list_axes[0].imshow(np.array(orig_img)) 
    list_axes[1].imshow(np.array(output_img)) 
    list_axes[0].set_title(orig_title, fontsize=15) 
    list_axes[1].set_title(output_title, fontsize=15) 
    fig.subplots_adjust(wspace=0.01 if is_horizontal else 0.00, hspace=0.01 if is_horizontal else 0.1) 
    fig.tight_layout() 
    return fig 

fig = visualize_results(original_image, images[0])

import gradio as gr 

def generate_from_text(text_promt, input_image, neg_prompt, seed, num_steps, style_strength_ratio): 
    """ 
    Helper function for generating result image from prompt text 

    Parameters: 
        text_promt (String): positive prompt 
        input_image (Image.Image): original image 
        neg_prompt (String): negative prompt 
        seed (Int): seed for random generator state initialization 
        num_steps (Int): number of sampling steps 
        style_strength_ratio (Int): the percentage of step when merging the ID embedding to text embedding 

    Returns: 
        result (Image.Image): generation result 
    """ 
    start_merge_step = int(float(style_strength_ratio) / 100 * num_steps) 
    if start_merge_step > 30: 
        start_merge_step = 30 
    result = pipe( 
        text_promt, 
        input_id_images=input_image, 
        negative_prompt=neg_prompt, 
        num_inference_steps=num_steps, 
        num_images_per_prompt=1, 

        start_merge_step=start_merge_step, 
        generator=torch.Generator().manual_seed(seed), 
        height=1024, 
        width=1024, 
    ).images[0] 

    return result 

with gr.Blocks() as demo: 
    with gr.Column(): 
        with gr.Row(): 
            input_image = gr.Image(label="Your image", sources=["upload"], type="pil") 
            output_image = gr.Image(label="Generated Images", type="pil") 
        positive_input = gr.Textbox(label=f"Text prompt, Trigger words is '{trigger_word}'") 
        neg_input = gr.Textbox(label="Negative prompt") 
        with gr.Row(): 
            seed_input = gr.Slider(0, 10_000_000, value=42, label="Seed") 
            steps_input = gr.Slider(label="Steps", value=10, minimum=5, maximum=50, step=1) 
            style_strength_ratio_input = gr.Slider(label="Style strength ratio", value=20, minimum=5, maximum=100, step=5) 
            btn = gr.Button() 
        btn.click( 
            generate_from_text, 
            [ 
                positive_input, 
                input_image, 
                neg_input, 
                seed_input, 
                steps_input, 
                style_strength_ratio_input, 
            ], 
            output_image, 
        ) 
        gr.Examples( 
            [ 
                [prompt, negative_prompt], 
                [ 
                    "A woman img wearing a Christmas hat", 
                    negative_prompt, 
                ], 
                [ 
                    "A man img in a helmet and vest riding a motorcycle", 
                    negative_prompt, 
                ], 
                [ 
                    "photo of a middle-aged man img sitting on a plush leather couch, and watching television show", 
                    negative_prompt, 
                ], 
                [ 
                    "photo of a skilled doctor img in a pristine white lab coat enjoying a delicious meal in a sophisticated dining room", 
                    negative_prompt, 
                ], 
                [ 
                    "photo of superman img flying through a vibrant sunset sky, with his cape billowing in the wind", 
                    negative_prompt, 
                ], 
            ], 
            [positive_input, neg_input], 
        ) 

demo.queue().launch() 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/

demo.close()