Spaces:
Runtime error
Runtime error
File size: 3,295 Bytes
f949b3f |
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 |
# import argparse
import sys
from pathlib import Path
from pytorch_lightning.cli import LightningCLI
from PIL import Image
# For streaming
import yaml
from copy import deepcopy
from typing import List, Optional
from jsonargparse.typing import restricted_string_type
# --------------------------------------
# ----------- For Streaming ------------
# --------------------------------------
class CustomCLI(LightningCLI):
def add_arguments_to_parser(self, parser):
parser.add_argument("--result_fol", type=Path,
help="Set the path to the result folder", default="results")
parser.add_argument("--exp_name", type=str, help="Experiment name")
parser.add_argument("--run_name", type=str,
help="Current run name")
parser.add_argument("--prompts", type=Optional[List[str]])
parser.add_argument("--scale_lr", type=bool,
help="Scale lr", default=False)
CodeType = restricted_string_type(
'CodeType', '(medium)|(high)|(highest)')
parser.add_argument("--matmul_precision", type=CodeType)
parser.add_argument("--ckpt", type=Path,)
parser.add_argument("--n_predictions", type=int)
return parser
def remove_value(dictionary, x):
for key, value in list(dictionary.items()):
if key == x:
del dictionary[key]
elif isinstance(value, dict):
remove_value(value, x)
return dictionary
def legacy_transformation(cfg: yaml):
cfg = deepcopy(cfg)
cfg["trainer"]["devices"] = "1"
cfg["trainer"]['num_nodes'] = 1
if not "class_path" in cfg["model"]["inference_params"]:
cfg["model"]["inference_params"] = {
"class_path": "t2v_enhanced.model.pl_module_params.InferenceParams", "init_args": cfg["model"]["inference_params"]}
return cfg
# ---------------------------------------------
# ----------- For enhancement -----------
# ---------------------------------------------
def add_margin(pil_img, top, right, bottom, left, color):
width, height = pil_img.size
new_width = width + right + left
new_height = height + top + bottom
result = Image.new(pil_img.mode, (new_width, new_height), color)
result.paste(pil_img, (left, top))
return result
def resize_to_fit(image, size):
W, H = size
w, h = image.size
if H / h > W / w:
H_ = int(h * W / w)
W_ = W
else:
W_ = int(w * H / h)
H_ = H
return image.resize((W_, H_))
def pad_to_fit(image, size):
W, H = size
w, h = image.size
pad_h = (H - h) // 2
pad_w = (W - w) // 2
return add_margin(image, pad_h, pad_w, pad_h, pad_w, (0, 0, 0))
def resize_and_keep(pil_img):
myheight = 576
hpercent = (myheight/float(pil_img.size[1]))
wsize = int((float(pil_img.size[0])*float(hpercent)))
pil_img = pil_img.resize((wsize, myheight))
return pil_img
def center_crop(pil_img):
width, height = pil_img.size
new_width = 576
new_height = 576
left = (width - new_width)/2
top = (height - new_height)/2
right = (width + new_width)/2
bottom = (height + new_height)/2
# Crop the center of the image
pil_img = pil_img.crop((left, top, right, bottom))
return pil_img |