129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Qwen-Image OpenAI-compatible image generation client.
|
|
|
|
Usage:
|
|
python openai_chat_client.py --prompt "A beautiful landscape" --output output.png
|
|
python openai_chat_client.py --prompt "A sunset" --height 1024 --width 1024 --steps 50 --seed 42
|
|
"""
|
|
|
|
import argparse
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
|
|
def generate_image(
|
|
prompt: str,
|
|
server_url: str = "http://localhost:8091",
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
steps: int | None = None,
|
|
true_cfg_scale: float | None = None,
|
|
seed: int | None = None,
|
|
negative_prompt: str | None = None,
|
|
num_outputs_per_prompt: int = 1,
|
|
) -> bytes | None:
|
|
"""Generate an image using the images generation API.
|
|
|
|
Args:
|
|
prompt: Text description of the image
|
|
server_url: Server URL
|
|
height: Image height in pixels
|
|
width: Image width in pixels
|
|
steps: Number of diffusion steps
|
|
true_cfg_scale: Qwen-Image CFG scale
|
|
seed: Random seed
|
|
negative_prompt: Negative prompt
|
|
num_outputs_per_prompt: Number of images to generate
|
|
|
|
Returns:
|
|
Image bytes or None if failed
|
|
"""
|
|
payload: dict[str, object] = {
|
|
"prompt": prompt,
|
|
"response_format": "b64_json",
|
|
"n": num_outputs_per_prompt,
|
|
}
|
|
|
|
if width is not None and height is not None:
|
|
payload["size"] = f"{width}x{height}"
|
|
elif width is not None:
|
|
payload["size"] = f"{width}x{width}"
|
|
elif height is not None:
|
|
payload["size"] = f"{height}x{height}"
|
|
|
|
if steps is not None:
|
|
payload["num_inference_steps"] = steps
|
|
if true_cfg_scale is not None:
|
|
payload["true_cfg_scale"] = true_cfg_scale
|
|
if negative_prompt:
|
|
payload["negative_prompt"] = negative_prompt
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
|
|
try:
|
|
response = requests.post(
|
|
f"{server_url}/v1/images/generations",
|
|
headers={"Content-Type": "application/json"},
|
|
json=payload,
|
|
timeout=300,
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
items = data.get("data")
|
|
if isinstance(items, list) and items:
|
|
first = items[0].get("b64_json") if isinstance(items[0], dict) else None
|
|
if isinstance(first, str):
|
|
return base64.b64decode(first)
|
|
|
|
print(f"Unexpected response format: {data}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
return None
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Qwen-Image chat client")
|
|
parser.add_argument("--prompt", "-p", default="a cup of coffee on the table", help="Text prompt")
|
|
parser.add_argument("--output", "-o", default="qwen_image_output.png", help="Output file")
|
|
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
|
|
parser.add_argument("--height", type=int, default=1024, help="Image height")
|
|
parser.add_argument("--width", type=int, default=1024, help="Image width")
|
|
parser.add_argument("--steps", type=int, default=50, help="Inference steps")
|
|
parser.add_argument("--cfg-scale", type=float, default=4.0, help="True CFG scale")
|
|
parser.add_argument("--seed", type=int, default=0, help="Random seed")
|
|
parser.add_argument("--negative", help="Negative prompt")
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"Generating image for: {args.prompt}")
|
|
|
|
image_bytes = generate_image(
|
|
prompt=args.prompt,
|
|
server_url=args.server,
|
|
height=args.height,
|
|
width=args.width,
|
|
steps=args.steps,
|
|
true_cfg_scale=args.cfg_scale,
|
|
seed=args.seed,
|
|
negative_prompt=args.negative,
|
|
)
|
|
|
|
if image_bytes:
|
|
output_path = Path(args.output)
|
|
output_path.write_bytes(image_bytes)
|
|
print(f"Image saved to: {output_path}")
|
|
print(f"Size: {len(image_bytes) / 1024:.1f} KB")
|
|
else:
|
|
print("Failed to generate image")
|
|
exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|