forked from Tencent/HunyuanDiT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sample_t2i.py
72 lines (59 loc) · 2.2 KB
/
sample_t2i.py
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
from pathlib import Path
from loguru import logger
from mllm.dialoggen_demo import DialogGen
from hydit.config import get_args
from hydit.inference import End2End
def inferencer():
args = get_args()
models_root_path = Path(args.model_root)
if not models_root_path.exists():
raise ValueError(f"`models_root` not exists: {models_root_path}")
# Load models
gen = End2End(args, models_root_path)
# Try to enhance prompt
if args.enhance:
logger.info("Loading DialogGen model (for prompt enhancement)...")
enhancer = DialogGen(str(models_root_path / "dialoggen"), args.load_4bit)
logger.info("DialogGen model loaded.")
else:
enhancer = None
return args, gen, enhancer
if __name__ == "__main__":
args, gen, enhancer = inferencer()
if enhancer:
logger.info("Prompt Enhancement...")
success, enhanced_prompt = enhancer(args.prompt)
if not success:
logger.info("Sorry, the prompt is not compliant, refuse to draw.")
exit()
logger.info(f"Enhanced prompt: {enhanced_prompt}")
else:
enhanced_prompt = None
# Run inference
logger.info("Generating images...")
height, width = args.image_size
results = gen.predict(args.prompt,
height=height,
width=width,
seed=args.seed,
enhanced_prompt=enhanced_prompt,
negative_prompt=args.negative,
infer_steps=args.infer_steps,
guidance_scale=args.cfg_scale,
batch_size=args.batch_size,
src_size_cond=args.size_cond,
)
images = results['images']
# Save images
save_dir = Path('results')
save_dir.mkdir(exist_ok=True)
# Find the first available index
all_files = list(save_dir.glob('*.png'))
if all_files:
start = max([int(f.stem) for f in all_files]) + 1
else:
start = 0
for idx, pil_img in enumerate(images):
save_path = save_dir / f"{idx + start}.png"
pil_img.save(save_path)
logger.info(f"Save to {save_path}")