|
1 | 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | +"""Run GLM-Image image-to-image (edit) generation through FastVideo. |
2 | 3 |
|
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 |
4 | 15 |
|
5 | 16 | from PIL import Image |
6 | 17 |
|
7 | 18 | 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 | +) |
8 | 29 |
|
9 | | -OUTPUT_PATH = "image_output" |
10 | | -CONDITION_IMAGE = "assets/images/couple.jpg" |
11 | 30 |
|
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() |
13 | 62 |
|
14 | 63 |
|
15 | 64 | 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() |
21 | 66 |
|
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"), |
35 | 85 | ) |
36 | 86 |
|
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] |
42 | 110 |
|
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() |
44 | 117 |
|
45 | 118 |
|
46 | 119 | if __name__ == "__main__": |
|
0 commit comments