PixArt-α: OpenVINO を使用したフォトリアリスティックなテキストから画像への合成の拡散トランスフォーマーの高速トレーニング#

%pip install -q "diffusers>=0.14.0" sentencepiece "datasets>=2.14.6" "transformers>=4.25.1" "gradio>=4.19" "torch>=2.1" Pillow opencv-python --extra-index-url https://download.pytorch.org/whl/cpu 
%pip install --pre -Uq openvino --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly

import torch 
from diffusers import PixArtAlphaPipeline 

pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-LCM-XL-2-1024-MS", use_safetensors=True) 

prompt = "A small cactus with a happy face in the Sahara desert." 
generator = torch.Generator().manual_seed(42) 

image = pipe(prompt, guidance_scale=0.0, num_inference_steps=4, 
generator=generator).images[0]

image

from pathlib import Path 

import numpy as np 
import torch 

import openvino as ov 

def convert(model: torch.nn.Module, xml_path: str, example_input): 
    xml_path = Path(xml_path) 
    if not xml_path.exists(): 
        xml_path.parent.mkdir(parents=True, exist_ok=True) 
        model.eval() 
        with torch.no_grad(): 
            converted_model = ov.convert_model(model, example_input=example_input) ov.save_model(converted_model, xml_path) 

        # メモリーをクリーンアップ 
        torch._C._jit_clear_class_registry() 
        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() 
        torch.jit._state._clear_class_state()

MODEL_DIR = Path("model") 

TEXT_ENCODER_PATH = MODEL_DIR / "text_encoder.xml" 
TRANSFORMER_OV_PATH = MODEL_DIR / "transformer_ir.xml" 
VAE_DECODER_PATH = MODEL_DIR / "vae_decoder.xml"

import ipywidgets as widgets 

core = ov.Core() 
device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

compiled_model = core.compile_model(TRANSFORMER_OV_PATH) 
compiled_vae = core.compile_model(VAE_DECODER_PATH) 
compiled_text_encoder = core.compile_model(TEXT_ENCODER_PATH)

from collections import namedtuple 

EncoderOutput = namedtuple("EncoderOutput", "last_hidden_state") 

class TextEncoderWrapper(torch.nn.Module): 
    def __init__(self, text_encoder, dtype): 
        super().__init__() 
        self.text_encoder = text_encoder 
        self.dtype = dtype 

    def forward(self, input_ids=None, attention_mask=None): 
        inputs = { 
            "input_ids": input_ids, 
            "attention_mask": attention_mask, 
        } 
        last_hidden_state = self.text_encoder(inputs)[0] 
        return EncoderOutput(torch.from_numpy(last_hidden_state))

class TransformerWrapper(torch.nn.Module): 
    def __init__(self, transformer, config): 
        super().__init__() 
        self.transformer = transformer 
        self.config = config 

    def forward( 
        self, 
        hidden_states=None, 
        timestep=None, 
        encoder_hidden_states=None, 
        encoder_attention_mask=None, 
        resolution=None, 
        aspect_ratio=None, 
        added_cond_kwargs=None, 
        **kwargs 
    ): 

        inputs = { 
            "hidden_states": hidden_states, 
            "timestep": timestep, 
            "encoder_hidden_states": encoder_hidden_states, 
            "encoder_attention_mask": encoder_attention_mask, 
        } 
        resolution = added_cond_kwargs["resolution"] 
        aspect_ratio = added_cond_kwargs["aspect_ratio"] 
        if resolution is not None: 
            inputs["resolution"] = resolution 
            inputs["aspect_ratio"] = aspect_ratio 
        outputs = self.transformer(inputs)[0] 

        return [torch.from_numpy(outputs)]

class VAEWrapper(torch.nn.Module): 
    def __init__(self, vae, config): 
        super().__init__() 
        self.vae = vae 
        self.config = config 

    def decode(self, latents=None, **kwargs): 
        inputs = { 
            "latents": latents, 
        } 

        outs = self.vae(inputs) 
        outs = namedtuple("VAE", "sample")(torch.from_numpy(outs[0])) 

    return outs

pipe.__dict__["_internal_dict"]["_execution_device"] = pipe._execution_device # これはパイプラインで発生する可能性のある問題を回避 

pipe.register_modules( 
    text_encoder=TextEncoderWrapper(compiled_text_encoder, pipe.text_encoder.dtype), 
    transformer=TransformerWrapper(compiled_model, pipe.transformer.config), 
    vae=VAEWrapper(compiled_vae, pipe.vae.config), 
)

generator = torch.Generator().manual_seed(42) 

image = pipe(prompt=prompt, guidance_scale=0.0, num_inference_steps=4, 
generator=generator).images[0]

image

import gradio as gr 

def generate(prompt, seed, negative_prompt, num_inference_steps): 
    generator = torch.Generator().manual_seed(seed) 
    image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, generator=generator, guidance_scale=0.0).images[0] 
    return image

demo = gr.Interface( 
    generate, 
    [ 
        gr.Textbox(label="Caption"), 
        gr.Slider(0, np.iinfo(np.int32).max, label="Seed"), 
        gr.Textbox(label="Negative prompt"), 
        gr.Slider(2, 20, step=1, label="Number of inference steps", value=4), 
    ], 
    "image", 
    examples=[ 
        ["A small cactus with a happy face in the Sahara desert.", 42], 
        ["an astronaut sitting in a diner, eating fries, cinematic, analog film", 42], 
        [ "Pirate ship trapped in a cosmic maelstrom nebula, rendered in cosmic beach whirlpool engine, volumetric lighting, spectacular, ambient lights, light pollution, cinematic atmosphere, art nouveau style, illustration art artwork by SenseiJaye, intricate detail.", 0, 
        ], 
        ["professional portrait photo of an anthropomorphic cat wearing fancy gentleman hat and jacket walking in autumn forest.", 0], 
    ], 
    allow_flagging="never", 
) 
try: 
    demo.queue().launch(debug=False) 
except Exception: 
    demo.queue().launch(debug=False, share=True) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細については、ドキュメントをご覧ください://gradio.app/docs/