Skip to content

Commit ee62319

Browse files
H1yori233SolitaryThinker
authored andcommitted
fix baed on review
1 parent 1e2f2bf commit ee62319

4 files changed

Lines changed: 235 additions & 220 deletions

File tree

examples/inference/basic/basic_glm_image.py

Lines changed: 88 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,106 @@
11
# SPDX-License-Identifier: Apache-2.0
2+
"""Run GLM-Image text-to-image generation through FastVideo.
23
3-
import os
4+
User story:
5+
"I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
6+
text-to-image generation command, saved as a PNG."
7+
"""
8+
import argparse
9+
from pathlib import Path
410

511
from PIL import Image
612

713
from fastvideo import VideoGenerator
14+
from fastvideo.api import (
15+
EngineConfig,
16+
GenerationRequest,
17+
GeneratorConfig,
18+
OutputConfig,
19+
ParallelismConfig,
20+
PipelineSelection,
21+
SamplingConfig,
22+
)
823

9-
OUTPUT_PATH = "image_output"
1024

11-
PROMPT = (
12-
"A beautiful landscape photography with rolling hills, "
13-
"a winding river, and a vibrant sunset in the background. "
14-
"Warm golden light, photorealistic style."
15-
)
25+
def parse_args() -> argparse.Namespace:
26+
parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
27+
parser.add_argument(
28+
"--model-path",
29+
default="zai-org/GLM-Image",
30+
help="HF id or local diffusers-format GLM-Image weights directory.",
31+
)
32+
parser.add_argument(
33+
"--output",
34+
default="image_output/landscape.png",
35+
help="Output PNG path.",
36+
)
37+
parser.add_argument(
38+
"--prompt",
39+
default=("A beautiful landscape photography with rolling hills, "
40+
"a winding river, and a vibrant sunset in the background. "
41+
"Warm golden light, photorealistic style."),
42+
help="Text prompt.",
43+
)
44+
parser.add_argument("--height", type=int, default=1024)
45+
parser.add_argument("--width", type=int, default=1024)
46+
parser.add_argument("--steps", type=int, default=50)
47+
parser.add_argument("--guidance-scale", type=float, default=1.5)
48+
parser.add_argument("--seed", type=int, default=1024)
49+
parser.add_argument("--num-gpus", type=int, default=1)
50+
parser.add_argument("--tp-size", type=int, default=None)
51+
parser.add_argument("--sp-size", type=int, default=None)
52+
return parser.parse_args()
1653

1754

1855
def main() -> None:
19-
generator = VideoGenerator.from_pretrained(
20-
"zai-org/GLM-Image",
21-
num_gpus=1,
22-
trust_remote_code=True,
23-
)
56+
args = parse_args()
2457

25-
os.makedirs(OUTPUT_PATH, exist_ok=True)
26-
27-
result = generator.generate_video(
28-
prompt=PROMPT,
29-
output_path=OUTPUT_PATH,
30-
save_video=False,
31-
return_frames=True,
32-
height=1024,
33-
width=1024,
34-
num_inference_steps=50,
35-
guidance_scale=1.5,
58+
output = Path(args.output)
59+
output.parent.mkdir(parents=True, exist_ok=True)
60+
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
61+
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
62+
63+
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
64+
# pipeline class come from the model's registered defaults — don't override.
65+
generator_config = GeneratorConfig(
66+
model_path=args.model_path,
67+
trust_remote_code=True,
68+
engine=EngineConfig(
69+
num_gpus=args.num_gpus,
70+
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
71+
),
72+
pipeline=PipelineSelection(workload_type="t2i"),
3673
)
3774

38-
frames = result.get("frames") if isinstance(result, dict) else None
39-
if frames:
40-
img = Image.fromarray(frames[0])
41-
img.save(os.path.join(OUTPUT_PATH, "landscape.png"))
42-
print(f"Saved image to {OUTPUT_PATH}/landscape.png")
75+
generator = VideoGenerator.from_config(generator_config)
76+
try:
77+
request = GenerationRequest(
78+
prompt=args.prompt,
79+
sampling=SamplingConfig(
80+
height=args.height,
81+
width=args.width,
82+
num_frames=1,
83+
fps=1,
84+
num_inference_steps=args.steps,
85+
guidance_scale=args.guidance_scale,
86+
seed=args.seed,
87+
),
88+
output=OutputConfig(
89+
output_path=str(output.parent),
90+
save_video=False,
91+
return_frames=True,
92+
),
93+
)
94+
result = generator.generate(request)
95+
if isinstance(result, list):
96+
result = result[0]
4397

44-
generator.shutdown()
98+
frames = result.frames
99+
if frames is not None and len(frames):
100+
Image.fromarray(frames[0]).save(output)
101+
print(f"Saved image to {output}")
102+
finally:
103+
generator.shutdown()
45104

46105

47106
if __name__ == "__main__":

examples/inference/basic/edit_glm_image.py

Lines changed: 101 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,46 +1,119 @@
11
# SPDX-License-Identifier: Apache-2.0
2+
"""Run GLM-Image image-to-image (edit) generation through FastVideo.
23
3-
import os
4+
User story:
5+
"I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
6+
want a minimal edit command (text + image -> edited image), saved as a PNG."
7+
8+
GLM-Image is a single unified pipeline: passing a condition image switches it
9+
from text-to-image to the edit path (the condition enters the DiT via a KV-cache
10+
write pass), so the generator config is identical to `basic_glm_image.py` — the
11+
`inputs.pil_image` on the request is what selects the edit mode.
12+
"""
13+
import argparse
14+
from pathlib import Path
415

516
from PIL import Image
617

718
from fastvideo import VideoGenerator
19+
from fastvideo.api import (
20+
EngineConfig,
21+
GenerationRequest,
22+
GeneratorConfig,
23+
InputConfig,
24+
OutputConfig,
25+
ParallelismConfig,
26+
PipelineSelection,
27+
SamplingConfig,
28+
)
829

9-
OUTPUT_PATH = "image_output"
10-
CONDITION_IMAGE = "assets/images/couple.jpg"
1130

12-
PROMPT = "Change the background to a snowy mountain landscape at golden hour."
31+
def parse_args() -> argparse.Namespace:
32+
parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
33+
parser.add_argument(
34+
"--model-path",
35+
default="zai-org/GLM-Image",
36+
help="HF id or local diffusers-format GLM-Image weights directory.",
37+
)
38+
parser.add_argument(
39+
"--image",
40+
default="assets/images/couple.jpg",
41+
help="Condition image to edit.",
42+
)
43+
parser.add_argument(
44+
"--output",
45+
default="image_output/edited.png",
46+
help="Output PNG path.",
47+
)
48+
parser.add_argument(
49+
"--prompt",
50+
default="Change the background to a snowy mountain landscape at golden hour.",
51+
help="Edit instruction.",
52+
)
53+
parser.add_argument("--height", type=int, default=1024)
54+
parser.add_argument("--width", type=int, default=1024)
55+
parser.add_argument("--steps", type=int, default=50)
56+
parser.add_argument("--guidance-scale", type=float, default=1.5)
57+
parser.add_argument("--seed", type=int, default=1024)
58+
parser.add_argument("--num-gpus", type=int, default=1)
59+
parser.add_argument("--tp-size", type=int, default=None)
60+
parser.add_argument("--sp-size", type=int, default=None)
61+
return parser.parse_args()
1362

1463

1564
def main() -> None:
16-
generator = VideoGenerator.from_pretrained(
17-
"zai-org/GLM-Image",
18-
num_gpus=1,
19-
trust_remote_code=True,
20-
)
65+
args = parse_args()
2166

22-
os.makedirs(OUTPUT_PATH, exist_ok=True)
23-
condition = Image.open(CONDITION_IMAGE).convert("RGB")
24-
25-
result = generator.generate_video(
26-
prompt=PROMPT,
27-
pil_image=condition,
28-
output_path=OUTPUT_PATH,
29-
save_video=False,
30-
return_frames=True,
31-
height=1024,
32-
width=1024,
33-
num_inference_steps=50,
34-
guidance_scale=1.5,
67+
output = Path(args.output)
68+
output.parent.mkdir(parents=True, exist_ok=True)
69+
condition = Image.open(args.image).convert("RGB")
70+
tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
71+
sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
72+
73+
# GLM-Image needs trust_remote_code for its AR encoder; offload and the
74+
# pipeline class come from the model's registered defaults — don't override.
75+
# The pipeline is registered as t2i; passing inputs.pil_image below switches
76+
# it to the edit path.
77+
generator_config = GeneratorConfig(
78+
model_path=args.model_path,
79+
trust_remote_code=True,
80+
engine=EngineConfig(
81+
num_gpus=args.num_gpus,
82+
parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
83+
),
84+
pipeline=PipelineSelection(workload_type="t2i"),
3585
)
3686

37-
frames = result.get("frames") if isinstance(result, dict) else None
38-
if frames:
39-
img = Image.fromarray(frames[0])
40-
img.save(os.path.join(OUTPUT_PATH, "edited.png"))
41-
print(f"Saved image to {OUTPUT_PATH}/edited.png")
87+
generator = VideoGenerator.from_config(generator_config)
88+
try:
89+
request = GenerationRequest(
90+
prompt=args.prompt,
91+
inputs=InputConfig(pil_image=condition),
92+
sampling=SamplingConfig(
93+
height=args.height,
94+
width=args.width,
95+
num_frames=1,
96+
fps=1,
97+
num_inference_steps=args.steps,
98+
guidance_scale=args.guidance_scale,
99+
seed=args.seed,
100+
),
101+
output=OutputConfig(
102+
output_path=str(output.parent),
103+
save_video=False,
104+
return_frames=True,
105+
),
106+
)
107+
result = generator.generate(request)
108+
if isinstance(result, list):
109+
result = result[0]
42110

43-
generator.shutdown()
111+
frames = result.frames
112+
if frames is not None and len(frames):
113+
Image.fromarray(frames[0]).save(output)
114+
print(f"Saved image to {output}")
115+
finally:
116+
generator.shutdown()
44117

45118

46119
if __name__ == "__main__":

fastvideo/pipelines/basic/glm_image/stages/denoising.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Ve
2929
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
3030
device = get_local_torch_device()
3131
dtype = torch.bfloat16
32-
guidance_scale = getattr(batch, "guidance_scale", 1.5)
32+
guidance_scale = batch.guidance_scale
3333
do_cfg = guidance_scale > 1.0
3434

3535
latents = getattr(batch, "latent", getattr(batch, "latents", None))

0 commit comments

Comments
 (0)