File size: 3,644 Bytes
bf63092 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import os
import torch
import numpy as np
import cv2
from PIL import Image
from argparse import ArgumentParser
from model import DiffMorpherPipeline
parser = ArgumentParser()
parser.add_argument(
"--model_path", type=str, default="stabilityai/stable-diffusion-2-1-base",
help="Pretrained model to use (default: %(default)s)"
)
parser.add_argument(
"--image_path_0", type=str, default="",
help="Path of the first image (default: %(default)s)")
parser.add_argument(
"--prompt_0", type=str, default="",
help="Prompt of the second image (default: %(default)s)")
parser.add_argument(
"--image_path_1", type=str, default="",
help="Path of the first image (default: %(default)s)")
parser.add_argument(
"--prompt_1", type=str, default="",
help="Prompt of the second image (default: %(default)s)")
parser.add_argument(
"--load_lora_path_0", type=str, default="",
help="Path of the lora directory of the first image (default: %(default)s)"
)
parser.add_argument(
"--load_lora_path_1", type=str, default="",
help="Path of the lora directory of the second image (default: %(default)s)"
)
parser.add_argument(
"--image_paths", type=str, nargs='*', default=[],
help="Path of the first image (default: %(default)s)")
parser.add_argument(
"--prompts", type=str, nargs='*', default=[],
help="Prompt of the second image (default: %(default)s)")
parser.add_argument(
"--output_path", type=str, default="./results",
help="Path of the output image (default: %(default)s)"
)
parser.add_argument(
"--save_lora_dir", type=str, default="./lora",
help="Path of the output lora directory (default: %(default)s)"
)
parser.add_argument(
"--load_lora_paths", type=str, nargs='*', default=[],
help="Path of the lora directory of the first image (default: %(default)s)"
)
parser.add_argument(
"--use_adain", action="store_true",
help="Use AdaIN (default: %(default)s)"
)
parser.add_argument(
"--use_reschedule", action="store_true",
help="Use reschedule sampling (default: %(default)s)"
)
parser.add_argument(
"--lamb", type=float, default=0.6,
help="Lambda for self-attention replacement (default: %(default)s)"
)
parser.add_argument(
"--fix_lora_value", type=float, default=None,
help="Fix lora value (default: LoRA Interp., not fixed)"
)
parser.add_argument(
"--save_inter", action="store_true",
help="Save intermediate results (default: %(default)s)"
)
parser.add_argument(
"--num_frames", type=int, default=16,
help="Number of frames to generate (default: %(default)s)"
)
parser.add_argument(
"--duration", type=int, default=100,
help="Duration of each frame (default: %(default)s ms)"
)
args = parser.parse_args()
os.makedirs(args.output_path, exist_ok=True)
pipeline = DiffMorpherPipeline.from_pretrained(
args.model_path, torch_dtype=torch.float32)
pipeline.to("cuda")
images = pipeline(
img_path_0=args.image_path_0,
img_path_1=args.image_path_1,
prompt_0=args.prompt_0,
prompt_1=args.prompt_1,
load_lora_path_0=args.load_lora_path_0,
load_lora_path_1=args.load_lora_path_1,
img_paths=args.image_paths,
prompts=args.prompts,
save_lora_dir=args.save_lora_dir,
load_lora_paths=args.load_lora_paths,
use_adain=args.use_adain,
use_reschedule=args.use_reschedule,
lamb=args.lamb,
output_path=args.output_path,
num_frames=args.num_frames,
fix_lora=args.fix_lora_value,
save_intermediates=args.save_inter,
)
images[0].save(f"{args.output_path}/output.gif", save_all=True,
append_images=images[1:], duration=args.duration, loop=0)
|