Mojo🔥でStable Diffusionを使うとPythonより1秒遅い

$Mojo$🔥で画像生成

$Python$でできることの殆どは$Mojo$🔥で書き直すことができる。以前$OpenCV$を使った際は、それだけで$Python$よりも高速になる場合があると分かった。

https://zenn.dev/amenaruya/articles/717015d9dfcec7

今回は画像生成を題材に、単に$Mojo$🔥で書き直すことで高速化されないものかと検証したが、そう旨い話はなかったようである。

簡単な実験

$Google$ $Colabolatory$で$Stable$ $Diffusion$を使う方法が簡単そうであった。ノートブックは$GitHub$にて共有する。

https://github.com/amenaruya/stable_diffusion_with_mojo

ランタイムはT4 GPUを選んだ。

Colabolatoryを使わない場合

自分の環境を使う場合

$Google$ $Colabolatory$を使わない場合についても、念のため簡単にまとめておく。但し$Mojo$🔥は$Windows$に対応していないことから、実行環境として$Windows$は必然に除外される。

インストール

torchは$GPU$の有無や$OS$の種類によって、適切にインストールしなければならない。

https://pytorch.org/get-started/locally/

コマンドの一例

その他は次の通りインストールできる。

# Stable Diffusion
pip install --upgrade diffusers[torch]

# Mojo
pip install modular

$Mojo$🔥はpipでインストールできるが、この方法は公式に推奨されているものではない。本来はpixiでインストールするのがよいとされている。

https://zenn.dev/amenaruya/articles/3b34658411863a

https://zenn.dev/amenaruya/articles/33723fd5465e71

$Python$

from diffusers import DiffusionPipeline
import torch
import time

pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline.to("cuda")

def __time__(func, *args, **kwargs):
    start = time.perf_counter()
    result = func(*args, **kwargs)
    end = time.perf_counter()
    elapsed = end - start
    print(f"{end - start:.6f} [s]")
    return result

prompt = input('prompt > ')
image = __time__(pipeline, prompt).images[0]

name = input('name > ')
image.save(f'{name}.png')

$Mojo$🔥

import time
import sys
from python import Python

fn main() raises:
    try:
        var torch = Python.import_module('torch')
        var diffusers = Python.import_module('diffusers')
        var pipeline = diffusers.DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16)
        pipeline.to("cuda")

        var prompt = input('prompt > ')

        var start = time.monotonic()
        var image = pipeline(prompt).images[0]
        var end = time.monotonic()
        var elapsed = (end - start) / 1_000_000_000
        print(String('{} [s]').format(elapsed))

        var name = input('file name > ')
        var file_name = String('{}.png').format(name)
        image.save(file_name)
        print(String('{} saved').format(file_name))
    except:
        sys.exit(1)

$Python$と$Mojo$🔥それぞれ$10$回ずつ試行し、その平均をとった。

	$Python$	$Mojo$🔥
$1$	$8.987383$	$8.749698725$
$2$	$9.248939$	$9.048349062$
$3$	$7.852054$	$8.657060594$
$4$	$7.934670$	$8.686166867$
$5$	$7.947472$	$9.687662112$
$6$	$7.984817$	$8.594926555$
$7$	$7.825276$	$9.673236295$
$8$	$7.981634$	$9.194612269$
$9$	$7.965697$	$9.547111223$
$10$	$7.819419$	$9.426520401$
平均	$8.154736$	$9.126534$

※平均はどちらも$Python$で算出したため、$Mojo$🔥のデータと桁数が揃っていないが、大した信憑性もないデータ相手に面倒なのでそのままにしている。

所感

$1$回目は僅かに$Mojo$🔥の方が速いが、単なる偶然だったらしい。目を疑うほど遅くなるというわけではなさそうだが、しかし確実に$Mojo$🔥の方が遅いようである。

本当は高速化という喜ばしい結果を期待していたのだが、以前に続きまたもや$Mojo$🔥のネガティブキャンペーンになってしまった。