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)