Stable Diffusion v2 と OpenVINO™ の無限ズーム#

%pip install -q "diffusers>=0.14.0" "transformers>=4.25.1" "gradio>=4.19" "openvino>=2024.2.0" "torch>=2.1" Pillow opencv-python "git+https://github.com/huggingface/optimum-intel.git" --extra-index-url https://download.pytorch.org/whl/cpu

import ipywidgets as widgets 
import openvino as ov 

core = ov.Core() 

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

device

from optimum.intel.openvino import OVStableDiffusionInpaintPipeline 
from pathlib import Path 

DEVICE = device.value 

MODEL_ID = "stabilityai/stable-diffusion-2-inpainting" 
MODEL_DIR = Path("sd2_inpainting") 

if not MODEL_DIR.exists(): 
    ov_pipe = OVStableDiffusionInpaintPipeline.from_pretrained(MODEL_ID, export=True, device=DEVICE, compile=False) 
    ov_pipe.save_pretrained(MODEL_DIR) 
else: 
    ov_pipe = OVStableDiffusionInpaintPipeline.from_pretrained(MODEL_DIR, device=DEVICE, compile=False) 

ov_pipe.compile()

from typing import List, Union 

import PIL 
import cv2 
from tqdm import trange 
import numpy as np 

def generate_video( 
    pipe, 
    prompt: Union[str, List[str]], 
    negative_prompt: Union[str, List[str]], 
    guidance_scale: float = 7.5, 
    num_inference_steps: int = 20, 
    num_frames: int = 20, 
    mask_width: int = 128, 
    seed: int = 9999, 
    zoom_in: bool = False, 
):     """ 
    Zoom video generation function 

    Parameters: 
        pipe (OVStableDiffusionInpaintingPipeline): inpainting pipeline. 
        prompt (str or List[str]): The prompt or prompts to guide the image generation. 
        negative_prompt (str or List[str]): The negative prompt or prompts to guide the image generation. 
        guidance_scale (float, *optional*, defaults to 7.5):
            Guidance scale as defined in Classifier-Free Diffusion Guidance(https://arxiv.org/abs/2207.12598). 
            guidance_scale is defined as `w` of equation 2.
            Higher guidance scale encourages to generate images that are closely linked to the text prompt, 
            usually at the expense of lower image quality. 
        num_inference_steps (int, *optional*, defaults to 50): The number of denoising steps for each frame. More denoising steps usually lead to a higher quality image at the expense of slower inference. 
        num_frames (int, *optional*, 20): number frames for video. 
        mask_width (int, *optional*, 128): size of border mask for inpainting on each step. 
        zoom_in (bool, *optional*, False): zoom mode Zoom In or Zoom Out.
     Returns: 
        output_path (str): Path where generated video loacated.
     """ 

    height = 512 
    width = height 

    current_image = PIL.Image.new(mode="RGBA", size=(height, width)) 
    mask_image = np.array(current_image)[:, :, 3] 
    mask_image = PIL.Image.fromarray(255 - mask_image).convert("RGB") 
    current_image = current_image.convert("RGB") 
    init_images = pipe( 
        prompt=prompt, 
        negative_prompt=negative_prompt, 
        image=current_image, 
        guidance_scale=guidance_scale, 
        mask_image=mask_image, 
        num_inference_steps=num_inference_steps, 
    ).images 

    image_grid(init_images, rows=1, cols=1) 

    num_outpainting_steps = num_frames 
    num_interpol_frames = 30 

    current_image = init_images[0] 
    all_frames = [] 
    all_frames.append(current_image) 
    for i in trange( 
        num_outpainting_steps, 
        desc=f"Generating {num_outpainting_steps} additional images...", 
    ): 
        prev_image_fix = current_image 

        prev_image = shrink_and_paste_on_blank(current_image, mask_width) 

        current_image = prev_image 

        # マスクを作成 (白い mask_width 幅のエッジを持つ黒い画像) 
        mask_image = np.array(current_image)[:, :, 3] 
        mask_image = PIL.Image.fromarray(255 - mask_image).convert("RGB") 

        # 修復ステップ 
        current_image = current_image.convert("RGB") 
        images = pipe( 
            prompt=prompt, 
            negative_prompt=negative_prompt, 
            image=current_image, 
            guidance_scale=guidance_scale,
            mask_image=mask_image, 
            num_inference_steps=num_inference_steps, 
        ).images 
        current_image = images[0] 
        current_image.paste(prev_image, mask=prev_image) 

        # 2 つの修復画像間の補間ステップ (= 連続ズームと切り取り) 
        for j in range(num_interpol_frames - 1): 
            interpol_image = current_image 
            interpol_width = round((1 - (1 - 2 * mask_width / height) ** (1 - (j + 1) / num_interpol_frames)) * height / 2) 
            interpol_image = interpol_image.crop( 
                ( 
                    interpol_width, 
                    interpol_width, 
                    width - interpol_width, 
                    height - interpol_width, 
                ) 
            ) 

            interpol_image = interpol_image.resize((height, width)) 

            # ズームによる画質の低下を避けるため、高解像度の前の画像を中央に貼り付け 
            interpol_width2 = round((1 - (height - 2 * mask_width) / (height - 2 * interpol_width)) / 2 * height) 
            prev_image_fix_crop = shrink_and_paste_on_blank(prev_image_fix, interpol_width2) 
            interpol_image.paste(prev_image_fix_crop, mask=prev_image_fix_crop) 
            all_frames.append(interpol_image) 
        all_frames.append(current_image) 
    video_file_name = f"infinite_zoom_{'in' if zoom_in else 'out'}" 
    fps = 30 
    save_path = video_file_name + ".mp4" 
    write_video(save_path, all_frames, fps, reversed_order=zoom_in) 
    return save_path

def shrink_and_paste_on_blank(current_image: PIL.Image.Image, mask_width: int): 
    """ 
    Decreases size of current_image by mask_width pixels from each side, 
    then adds a mask_width width transparent frame, 
    so that the image the function returns is the same size as the input.
    Parameters: 
        current_image (PIL.Image): input image to transform 
        mask_width (int): width in pixels to shrink from each side 
    Returns: 
        prev_image (PIL.Image): resized image with extended borders 
    """ 

    height = current_image.height 
    width = current_image.width 

    # mask_width で縮小 
    prev_image = current_image.resize((height - 2 * mask_width, width - 2 * mask_width)) 
    prev_image = prev_image.convert("RGBA") 
    prev_image = np.array(prev_image) 

    # 空白の不透明画像を作成 
    blank_image = np.array(current_image.convert("RGBA")) * 0 
    blank_image[:, :, 3] = 1 

    # 縮小して空白部分に貼り付け 
    blank_image[mask_width : height - mask_width, mask_width : width - mask_width, :]= prev_image 
    prev_image = PIL.Image.fromarray(blank_image) 

    return prev_image 

def image_grid(imgs: List[PIL.Image.Image], rows: int, cols: int): 
    """ 
    Insert images to grid 

    Parameters: 
        imgs (List[PIL.Image.Image]): list of images for making grid 
        rows (int): number of rows in grid 
        cols (int): number of columns in grid 
    Returns: 
        grid (PIL.Image): image with input images collage 
    """ 
    assert len(imgs) == rows * cols 

    w, h = imgs[0].size 
    grid = PIL.Image.new("RGB", size=(cols * w, rows * h)) 

    for i, img in enumerate(imgs): 
        grid.paste(img, box=(i % cols * w, i // cols * h)) 
    return grid 

def write_video( 
    file_path: str, 
    frames: List[PIL.Image.Image], 
    fps: float, 
    reversed_order: bool = True, 
    gif: bool = True, 
):     """ 
    Writes frames to an mp4 video file and optionaly to gif 

    Parameters: 
        file_path (str): Path to output video, must end with .mp4 
        frames (List of PIL.Image): list of frames 
        fps (float): Desired frame rate 
        reversed_order (bool): if order of images to be reversed (default = True) 
        gif (bool): save frames to gif format (default = True) 
    Returns: None 
    """ 
    if reversed_order: 
        frames.reverse() 

    w, h = frames[0].size 
    fourcc = cv2.VideoWriter_fourcc("m", "p", "4", "v") 
    writer = cv2.VideoWriter(file_path, fourcc, fps, (w, h)) 

    for frame in frames: 
        np_frame = np.array(frame.convert("RGB")) 
        cv_frame = cv2.cvtColor(np_frame, cv2.COLOR_RGB2BGR) 
        writer.write(cv_frame) 

    writer.release() 
    if gif: 
        frames[0].save( 
            file_path.replace(".mp4", ".gif"), 
            save_all=True, 
            append_images=frames[1:], 
            duratiobn=len(frames) / fps, 
            loop=0, 
        )

import gradio as gr 

def generate( 
    prompt, 
    negative_prompt, 
    seed, 
    steps, 
    frames, 
    edge_size, 
    zoom_in, 
    progress=gr.Progress(track_tqdm=True), 
): 
    np.random.seed(seed) 
    video_path = generate_video( 
        ov_pipe, 
        prompt, 
        negative_prompt, 
        num_inference_steps=steps, 
        num_frames=frames, 
        mask_width=edge_size, 
        zoom_in=zoom_in, 
    ) 
    np.random.seed(None) 

    return video_path.replace(".mp4", ".gif") 

gr.close_all() 
demo = gr.Interface( 
    generate, 
    [ 
        gr.Textbox( 
            "valley in the Alps at sunset, epic vista, beautiful landscape, 4k, 8k",
            label="Prompt", 
        ), 
        gr.Textbox("lurry, bad art, blurred, text, watermark", label="Negative prompt"), 
        gr.Slider(value=9999, label="Seed", step=1, maximum=10000000), 
        gr.Slider(value=20, label="Steps", minimum=1, maximum=50), 
        gr.Slider(value=3, label="Frames", minimum=1, maximum=50), 
        gr.Slider(value=128, label="Edge size", minimum=32, maximum=256), 
        gr.Checkbox(label="Zoom in"), 
    ], 
    "image", 
) 

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