File size: 3,989 Bytes
2c87bcd 74f59c5 2c87bcd 74f59c5 2c87bcd 74f59c5 1a625a7 2c87bcd 74f59c5 2c87bcd 74f59c5 2c87bcd 74f59c5 30864be 2c87bcd 74f59c5 2c87bcd 74f59c5 2c87bcd 74f59c5 2c87bcd 30864be 2c87bcd 74f59c5 30864be 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os
import sys
import subprocess
import base64
import io
import traceback
# Install latest diffusers (supports GLM-Image) + dependencies
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
"diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "transformers", "huggingface_hub"])
import torch
from flask import Flask, request, jsonify
print(f"[image] torch version: {torch.__version__}", flush=True)
print(f"[image] CUDA available: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
print(f"[image] GPU: {torch.cuda.get_device_name(0)}", flush=True)
print(f"[image] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)
app = Flask(__name__)
pipe = None
def load_model():
global pipe
from diffusers.pipelines.glm_image import GlmImagePipeline
print("[image] Loading GLM-Image pipeline...", flush=True)
pipe = GlmImagePipeline.from_pretrained(
"zai-org/GLM-Image",
torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
print("[image] Model loaded on CUDA!", 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(force=True)
# Support both HF format (inputs) and backend format (prompt)
prompt = data.get("inputs", "") or data.get("prompt", "")
params = data.get("parameters", {})
width = int(params.get("width", 1024))
height = int(params.get("height", 1024))
steps = int(params.get("num_inference_steps", 30))
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)
print(f"[image] Params: {width}x{height}, steps={steps}, guidance={guidance}", 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, size: {len(buf.getvalue())} bytes", flush=True)
# Return in both formats - HF format and backend format
return jsonify({
"image": img_b64,
"format": "png",
"success": True,
"image_base64": img_b64,
"model": "glm-image",
})
except Exception as e:
tb = traceback.format_exc()
print(f"[image] ERROR: {e}\n{tb}", flush=True)
return jsonify({"error": str(e), "traceback": tb, "success": False}), 500
# Also add /api/generate route for the backend's expected format
@app.route("/api/generate", methods=["POST"])
def api_generate():
"""Handle the backend's /api/generate format by forwarding to generate()"""
return generate()
# Also add /api/analyze and /api/vision-chat for the backend
@app.route("/api/analyze", methods=["POST"])
def api_analyze():
"""Image analysis - return a placeholder since we don't have a vision model here"""
data = request.get_json(force=True)
return jsonify({
"success": True,
"description": "Image analysis is not available on this endpoint.",
"model": "none",
})
@app.route("/api/vision-chat", methods=["POST"])
def api_vision_chat():
"""Vision chat - return a placeholder"""
data = request.get_json(force=True)
return jsonify({
"success": True,
"answer": "Vision chat is not available on this endpoint.",
"model": "none",
})
if __name__ == "__main__":
load_model()
app.run(host="0.0.0.0", port=8000)
|