76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
GLM-Image inference example using FastVideo.
|
|
|
|
This example demonstrates how to generate images using the GLM-Image model
|
|
through FastVideo's pipeline infrastructure.
|
|
|
|
Usage:
|
|
python examples/inference/basic/basic_glm_image.py
|
|
"""
|
|
|
|
import os
|
|
|
|
# Set environment variables for single-GPU distributed setup
|
|
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
|
os.environ.setdefault("MASTER_PORT", "29501")
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
|
|
from fastvideo.pipelines.basic.glm_image import GlmImagePipeline
|
|
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
|
|
|
OUTPUT_PATH = "image_output"
|
|
|
|
|
|
|
|
def main():
|
|
# Load the GLM-Image pipeline
|
|
pipe = GlmImagePipeline.from_pretrained(
|
|
"zai-org/GLM-Image",
|
|
num_gpus=1,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
# Create a batch with generation parameters
|
|
prompt = (
|
|
"A beautiful landscape photography with rolling hills, "
|
|
"a winding river, and a vibrant sunset in the background. "
|
|
"Warm golden light, photorealistic style."
|
|
)
|
|
|
|
batch = ForwardBatch(
|
|
data_type="image",
|
|
prompt=prompt,
|
|
negative_prompt="",
|
|
height=1024,
|
|
width=1024,
|
|
num_inference_steps=50,
|
|
guidance_scale=7.5,
|
|
guidance_rescale=0.7,
|
|
do_classifier_free_guidance=True,
|
|
num_frames=1,
|
|
seed=42,
|
|
)
|
|
|
|
# Generate the image
|
|
result = pipe.forward(batch, pipe.fastvideo_args)
|
|
|
|
# Output is in result.output as a tensor [B, C, T, H, W]
|
|
print(f"Generated image tensor shape: {result.output.shape}")
|
|
|
|
# Save the output image
|
|
import torch
|
|
from torchvision.utils import save_image
|
|
|
|
os.makedirs(OUTPUT_PATH, exist_ok=True)
|
|
# result.output is [B, C, 1, H, W] for images
|
|
image_tensor = result.output.squeeze(2) # [B, C, H, W]
|
|
save_path = os.path.join(OUTPUT_PATH, "output.png")
|
|
save_image(image_tensor, save_path)
|
|
print(f"Image saved to {save_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|