Segmind-VegaRT と OpenVINO による高解像度の画像生成#

%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu\ "torch>=2.1" transformers diffusers "git+https://github.com/huggingface/optimum-intel.git" "gradio>=4.19" "openvino>=2023.3.0" "peft==0.6.2"

import torch 
from diffusers import LCMScheduler, AutoPipelineForText2Image 
import gc 
from pathlib import Path 

model_id = "segmind/Segmind-Vega" 
adapter_id = "segmind/Segmind-VegaRT" 
pt_model_dir = Path("segmind-vegart") 

if not pt_model_dir.exists(): 
    pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16") 
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) 
    pipe.load_lora_weights(adapter_id) 
    pipe.fuse_lora()
 
    pipe.save_pretrained("segmind-vegart") 
    del pipe 
    gc.collect()

from pathlib import Path 

model_dir = Path("openvino-segmind-vegart") 
sdxl_model_id = "./segmind-vegart" 
tae_id = "madebyollin/taesdxl" 
skip_convert_model = model_dir.exists()

import torch 
import openvino as ov 
from diffusers import AutoencoderTiny 
import gc 

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

    def forward(self, sample): 
        return self.vae.encode(sample) 

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

    def forward(self, latent_sample): 
        return self.vae.decode(latent_sample) 

def convert_tiny_vae(model_id, output_path): 
    tiny_vae = AutoencoderTiny.from_pretrained(model_id) 
    tiny_vae.eval() 
    vae_encoder = VAEEncoder(tiny_vae) 
    ov_model = ov.convert_model(vae_encoder, example_input=torch.zeros((1, 3, 512, 512))) 
    ov.save_model(ov_model, output_path / "vae_encoder/openvino_model.xml") 
    tiny_vae.save_config(output_path / "vae_encoder") 
    vae_decoder = VAEDecoder(tiny_vae) 
    ov_model = ov.convert_model(vae_decoder, example_input=torch.zeros((1, 4, 64, 64))) 
    ov.save_model(ov_model, output_path / "vae_decoder/openvino_model.xml") 
    tiny_vae.save_config(output_path / "vae_decoder") 
    del tiny_vae 
    del ov_model 
    gc.collect() 

if not skip_convert_model:
    !optimum-cli export openvino --model $sdxl_model_id --task stable-diffusion-xl $model_dir --fp16 
    convert_tiny_vae(tae_id, model_dir)

to_quantize = widgets.Checkbox( 
    value=True, 
    description="Quantization", 
    disabled=False, 
) 

to_quantize

# Fetch `skip_kernel_extension` module 
import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py", 
) 
open("skip_kernel_extension.py", "w").write(r.text) 

int8_pipe = None 

if to_quantize.value and "GPU" in device.value: 
    to_quantize.value = False 
42
%load_ext skip_kernel_extension

quantized_model_present = UNET_INT8_OV_PATH.exists() 

use_quantized_model = widgets.Checkbox( 
    value=quantized_model_present, 
    description="Use quantized model", 
    disabled=not quantized_model_present, 
) 

use_quantized_model

import gradio as gr 

if use_quantized_model.value: 
    if not quantized_model_present: 
        raise RuntimeError("Quantized model not found.") 
    text2image_pipe = create_int8_pipe(model_dir, UNET_INT8_OV_PATH, device.value, core) 
else: 
    text2image_pipe = OVStableDiffusionXLPipeline.from_pretrained(model_dir, device=device.value) 

def generate_from_text(text, seed, num_steps, height, width): 
    set_seed(seed) 
    result = text2image_pipe( 
        text, 
        num_inference_steps=num_steps, 
        guidance_scale=1.0, 
        height=height, 
        width=width, 
    ).images[0] 
    return result 

with gr.Blocks() as demo: 
    with gr.Column(): 
        positive_input = gr.Textbox(label="Text prompt") 
        with gr.Row(): 
            seed_input = gr.Number(precision=0, label="Seed", value=42, minimum=0) 
            steps_input = gr.Slider(label="Steps", value=4, minimum=2, maximum=8, step=1) 
            height_input = gr.Slider(label="Height", value=512, minimum=256, maximum=1024, step=32) 
            width_input = gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=32) 
            btn = gr.Button() 
        out = gr.Image( 
            label=("Result (Quantized)" if use_quantized_model.value else "Result (Original)"), 
            type="pil", 
            width=512, 
        ) 
        btn.click( 
            generate_from_text, 
            [positive_input, seed_input, steps_input, height_input, width_input], 
            out, 
        ) 
        gr.Examples( 
            [ 
                ["cute cat", 999], 
                [ 
                    "underwater world coral reef, colorful jellyfish, 35mm, cinematic lighting, shallow depth of field, ultra quality, masterpiece, realistic", 
                    89, 
                ], 
                [ 
                    "a photo realistic happy white poodle dog ​​playing in the grass, extremely detailed, high res, 8k, masterpiece, dynamic angle", 
                    1569, 
                ], 
                [ 
                    "Astronaut on Mars watching sunset, best quality, cinematic effects,", 
                    65245, 
                ], 
                [ 
                    "Black and white street photography of a rainy night in New York, reflections on wet pavement", 
                    48199, 
                ], 
                [ 
                    "cinematic photo detailed closeup portraid of a Beautiful cyberpunk woman, robotic parts, cables, lights, text; , high quality photography, 3 point lighting, flash with softbox, 4k, Canon EOS R3, hdr, smooth, sharp focus, high resolution, award winning photo, 80mm, f2.8, bokeh .35mm photograph, film, bokeh, professional, 4k, highly detailed, high quality photography, 3 point lighting, flash with softbox, 4k, Canon EOS R3, hdr, smooth, sharp focus, high resolution, award winning photo, 80mm, f2.8, bokeh", 
                    48199, 
                ], 
            ], 
            [positive_input, seed_input], 
        ) 

# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/ 
# デモを共有するため公開リンクを作成したい場合は、share=True を追加 
try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(share=True, debug=False)