File size: 2,254 Bytes
2c87bcd 2ca5d8d 2c87bcd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import os
import sys
import subprocess
import base64
import io
import json
# Install diffusers (latest PyPI supports GLM-Image) + flask for serving
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", "diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "transformers"])
import torch
from flask import Flask, request, jsonify
from diffusers.pipelines.glm_image import GlmImagePipeline
app = Flask(__name__)
pipe = None
def load_model():
global pipe
print("[image] Loading GLM-Image pipeline...", flush=True)
pipe = GlmImagePipeline.from_pretrained(
"zai-org/GLM-Image",
torch_dtype=torch.bfloat16,
device_map="cuda",
enable_model_cpu_offload=True,
)
print("[image] Model loaded successfully!", flush=True)
@app.route("/health", methods=["GET"])
def health():
if pipe is not None:
return jsonify({"status": "healthy"}), 200
return jsonify({"status": "loading"}), 503
@app.route("/", methods=["POST"])
def generate():
try:
data = request.get_json()
prompt = data.get("inputs", "")
params = data.get("parameters", {})
width = int(params.get("width", 1024))
height = int(params.get("height", 1024))
steps = int(params.get("num_inference_steps", 50))
guidance = float(params.get("guidance_scale", 1.5))
# GLM-Image requires dimensions divisible by 32
width = (width // 32) * 32
height = (height // 32) * 32
print(f"[image] Generating: {prompt[:100]}...", flush=True)
image = pipe(
prompt=prompt,
height=height,
width=width,
num_inference_steps=steps,
guidance_scale=guidance,
).images[0]
buf = io.BytesIO()
image.save(buf, format="PNG")
img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
print(f"[image] Done, image size: {len(buf.getvalue())} bytes", flush=True)
return jsonify({"image": img_b64, "format": "png"})
except Exception as e:
print(f"[image] Error: {e}", flush=True)
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
load_model()
app.run(host="0.0.0.0", port=8000)
|