5037af1033
Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk> Signed-off-by: Andy Zhou <jzhoubc@connect.ust.hk> Signed-off-by: dongbo910220 <1275604947@qq.com> Signed-off-by: Andy Zhou <46011930+AndyZhou952@users.noreply.github.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Samit <285365963@qq.com>
125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
OpenAI-compatible chat client for diffusion LoRA inference.
|
|
|
|
Example:
|
|
python openai_chat_client.py \
|
|
--server http://localhost:8091 \
|
|
--prompt "A piece of cheesecake" \
|
|
--lora-path /path/to/lora_adapter \
|
|
--lora-name my_lora \
|
|
--lora-scale 1.0 \
|
|
--output output.png
|
|
"""
|
|
|
|
import argparse
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
|
|
def generate_image(
|
|
*,
|
|
prompt: str,
|
|
server_url: str,
|
|
height: int | None,
|
|
width: int | None,
|
|
num_inference_steps: int | None,
|
|
seed: int | None,
|
|
lora_name: str | None,
|
|
lora_path: str | None,
|
|
lora_scale: float | None,
|
|
lora_int_id: int | None,
|
|
) -> bytes | None:
|
|
messages = [{"role": "user", "content": prompt}]
|
|
|
|
extra_body: dict = {}
|
|
if height is not None:
|
|
extra_body["height"] = height
|
|
if width is not None:
|
|
extra_body["width"] = width
|
|
if num_inference_steps is not None:
|
|
extra_body["num_inference_steps"] = num_inference_steps
|
|
if seed is not None:
|
|
extra_body["seed"] = seed
|
|
|
|
if lora_path:
|
|
lora_body: dict = {
|
|
"local_path": lora_path,
|
|
"name": lora_name or Path(lora_path).stem,
|
|
}
|
|
if lora_scale is not None:
|
|
lora_body["scale"] = float(lora_scale)
|
|
if lora_int_id is not None:
|
|
lora_body["int_id"] = int(lora_int_id)
|
|
extra_body["lora"] = lora_body
|
|
|
|
payload = {"messages": messages}
|
|
if extra_body:
|
|
payload["extra_body"] = extra_body
|
|
|
|
response = requests.post(
|
|
f"{server_url}/v1/chat/completions",
|
|
headers={"Content-Type": "application/json"},
|
|
json=payload,
|
|
timeout=300,
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
content = data["choices"][0]["message"]["content"]
|
|
if isinstance(content, list) and content:
|
|
image_url = content[0].get("image_url", {}).get("url", "")
|
|
if image_url.startswith("data:image"):
|
|
_, b64_data = image_url.split(",", 1)
|
|
return base64.b64decode(b64_data)
|
|
|
|
raise RuntimeError(f"Unexpected response format: {content!r}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Diffusion LoRA OpenAI chat client")
|
|
parser.add_argument("--server", default="http://localhost:8091", help="Server URL")
|
|
parser.add_argument("--prompt", default="A piece of cheesecake", help="Text prompt")
|
|
parser.add_argument("--output", default="lora_online_output.png", help="Output image path")
|
|
|
|
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="num_inference_steps")
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
|
|
|
parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")
|
|
parser.add_argument("--lora-name", default=None, help="LoRA name (optional)")
|
|
parser.add_argument("--lora-scale", type=float, default=1.0, help="LoRA scale")
|
|
parser.add_argument(
|
|
"--lora-int-id",
|
|
type=int,
|
|
default=None,
|
|
help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
image_bytes = generate_image(
|
|
prompt=args.prompt,
|
|
server_url=args.server,
|
|
height=args.height,
|
|
width=args.width,
|
|
num_inference_steps=args.steps,
|
|
seed=args.seed,
|
|
lora_name=args.lora_name,
|
|
lora_path=args.lora_path,
|
|
lora_scale=args.lora_scale if args.lora_path else None,
|
|
lora_int_id=args.lora_int_id if args.lora_path else None,
|
|
)
|
|
|
|
out_path = Path(args.output)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
out_path.write_bytes(image_bytes)
|
|
print(f"Saved: {out_path} ({len(image_bytes) / 1024:.1f} KiB)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|