Upload server.py with huggingface_hub

server.py

@@ -0,0 +1,73 @@
+import os
+import sys
+import subprocess
+import base64
+import io
+import json
+
+# Install diffusers from source (supports GLM-Image)
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "git+https://github.com/huggingface/diffusers.git"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "accelerate", "sentencepiece", "protobuf"])
+
+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)