multi-view-diffusion / run_imagedream.py
ashawkey's picture
both worked... but imagedream quality is unmatched
3a2ea0a
raw
history blame
900 Bytes
import torch
import kiui
import numpy as np
import argparse
from mvdream.pipeline_mvdream import MVDreamPipeline
pipe = MVDreamPipeline.from_pretrained(
"./weights_imagedream", # local weights
# "ashawkey/mvdream-sd2.1-diffusers",
torch_dtype=torch.float16
)
pipe = pipe.to("cuda")
parser = argparse.ArgumentParser(description="ImageDream")
parser.add_argument("image", type=str, default='data/anya_rgba.png')
parser.add_argument("--prompt", type=str, default="")
args = parser.parse_args()
while True:
input_image = kiui.read_image(args.image, mode='float')
image = pipe(args.prompt, input_image)
grid = np.concatenate(
[
np.concatenate([image[0], image[2]], axis=0),
np.concatenate([image[1], image[3]], axis=0),
],
axis=1,
)
# kiui.vis.plot_image(grid)
kiui.write_image('test_imagedream.jpg', grid)
break