Stable Cascade と OpenVINO による画像生成#

%pip install -q "diffusers>=0.27.0" accelerate datasets gradio transformers "nncf>=2.10.0" "openvino>=2024.1.0" "torch>=2.1" --extra-index-url https://download.pytorch.org/whl/cpu

import torch 
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline 

prompt = "an image of a shiba inu, donning a spacesuit and helmet" 
negative_prompt = "" 

prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.float32) 
decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float32)

import ipywidgets as widgets 

run_original_inference = widgets.Checkbox( 
    value=False, 
    description="Run original inference", 
    disabled=False, 
) 

run_original_inference

if run_original_inference.value: 
    prior.to(torch.device("cpu")) 
    prior_output = prior( 
        prompt=prompt, 
        height=1024, 
        width=1024, 
        negative_prompt=negative_prompt, 
        guidance_scale=4.0, 
        num_images_per_prompt=1, 
        num_inference_steps=20, 
    ) 

    decoder_output = decoder( 
        image_embeddings=prior_output.
        image_embeddings, prompt=prompt, 
        negative_prompt=negative_prompt, 
        guidance_scale=0.0, 
        output_type="pil", 
        num_inference_steps=10, 
).images[0] 
display(decoder_output)

import gc 
from pathlib import Path 

import openvino as ov 
import nncf 

MODELS_DIR = Path("models") 

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

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

    gc.collect()

core = ov.Core() 

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

device

from collections import namedtuple 

BaseModelOutputWithPooling = namedtuple("BaseModelOutputWithPooling", ["text_embeds", "last_hidden_state", "hidden_states"]) 

class TextEncoderWrapper: 
    dtype = torch.float32 # accessed in the original workflow 

    def __init__(self, text_encoder_path, device): 
        self.text_encoder = core.compile_model(text_encoder_path, device.value) 

    def __call__(self, input_ids, attention_mask, output_hidden_states=True): 
        output = self.text_encoder({"input_ids": input_ids, "attention_mask": attention_mask}) 
        text_embeds = output[0] 
        last_hidden_state = output[1] 
        hidden_states = list(output.values())[1:] 
        return BaseModelOutputWithPooling(torch.from_numpy(text_embeds), torch.from_numpy(last_hidden_state), [torch.from_numpy(hs) for hs in hidden_states])

class PriorPriorWrapper: 
    def __init__(self, prior_path, device): 
        self.prior = core.compile_model(prior_path, device.value) 
        self.config = namedtuple("PriorWrapperConfig", ["clip_image_in_channels", "in_channels"])(768, 16) # 元のワークフローでアクセス 
        self.parameters = lambda: (torch.zeros(i, dtype=torch.float32) for i in range(1)) # 元のワークフローでアクセス 

    def __call__(self, sample, timestep_ratio, clip_text_pooled, clip_text=None, clip_img=None, **kwargs): 
        inputs = { 
            "sample": sample, 
            "timestep_ratio": timestep_ratio, 
            "clip_text_pooled": clip_text_pooled, 
            "clip_text": clip_text, 
            "clip_img": clip_img, 
        } 
        output = self.prior(inputs) 
        return [torch.from_numpy(output[0])]

class DecoderWrapper: 
    dtype = torch.float32 # 元のワークフローでアクセス 

    def __init__(self, decoder_path, device): 
        self.decoder = core.compile_model(decoder_path, device.value) 

    def __call__(self, sample, timestep_ratio, clip_text_pooled, effnet, **kwargs): 
        inputs = {"sample": sample, "timestep_ratio": timestep_ratio, "clip_text_pooled": clip_text_pooled, "effnet": effnet} 
        output = self.decoder(inputs) 
        return [torch.from_numpy(output[0])]

VqganOutput = namedtuple("VqganOutput", "sample") 

class VqganWrapper: 
    config = namedtuple("VqganWrapperConfig", "scale_factor")(0.3764) # 元のワークフローでアクセス 

    def __init__(self, vqgan_path, device): 
        self.vqgan = core.compile_model(vqgan_path, device.value) 

    def decode(self, h): 
        output = self.vqgan(h)[0] 
        output = torch.tensor(output) 
        return VqganOutput(output)

prior.text_encoder = TextEncoderWrapper(PRIOR_TEXT_ENCODER_OV_PATH, device) 
prior.prior = PriorPriorWrapper(PRIOR_PRIOR_MODEL_OV_PATH, device) 
decoder.decoder = DecoderWrapper(DECODER_DECODER_MODEL_OV_PATH, device) 
decoder.text_encoder = TextEncoderWrapper(DECODER_TEXT_ENCODER_MODEL_OV_PATH, device) 
decoder.vqgan = VqganWrapper(VQGAN_PATH, device)

prior_output = prior( 
    prompt=prompt, 
    height=1024, 
    width=1024, 
    negative_prompt=negative_prompt, 
    guidance_scale=4.0, 
    num_images_per_prompt=1, 
    num_inference_steps=20, 
) 

decoder_output = decoder( 
    image_embeddings=prior_output.image_embeddings, 
    prompt=prompt, 
    negative_prompt=negative_prompt, 
    guidance_scale=0.0, 
    output_type="pil", 
    num_inference_steps=10, 
).images[0] 
display(decoder_output)

def generate(prompt, negative_prompt, prior_guidance_scale, decoder_guidance_scale, seed): 
    generator = torch.Generator().manual_seed(seed) 
    prior_output = prior( 
        prompt=prompt, 
        height=1024, 
        width=1024, 
        negative_prompt=negative_prompt, 
        guidance_scale=prior_guidance_scale, 
        num_images_per_prompt=1, 
        num_inference_steps=20, 
        generator=generator, 
) 

decoder_output = decoder( 
    image_embeddings=prior_output.image_embeddings, 
    prompt=prompt, 
    negative_prompt=negative_prompt, 
    guidance_scale=decoder_guidance_scale, 
    output_type="pil", 
    num_inference_steps=10, 
    generator=generator, 
).images[0] 

return decoder_output

import gradio as gr 
import numpy as np 

demo = gr.Interface( 
    generate, 
    [ 
        gr.Textbox(label="Prompt"), 
        gr.Textbox(label="Negative prompt"), 
        gr.Slider( 
            0, 
            20, 
            step=1, 
            label="Prior guidance scale", 
            info="Higher guidance scale encourages to generate images that are closely " 
            "linked to the text `prompt`, usually at the expense of lower image quality.Applies to the prior pipeline", 
        ), 
        gr.Slider( 
            0, 
            20, 
            step=1, 
            label="Decoder guidance scale", 
            info="Higher guidance scale encourages to generate images that are closely " 
            "linked to the text `prompt`, usually at the expense of lower image quality.Applies to the decoder pipeline", 
        ), 
        gr.Slider(0, np.iinfo(np.int32).max, label="Seed", step=1), 
    ], 
    "image", 
    examples=[["An image of a shiba inu, donning a spacesuit and helmet", "", 4, 0, 0], ["An armchair in the shape of an avocado", "", 4, 0, 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') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/