Spaces:
Sleeping
Sleeping
| import io | |
| import logging | |
| import os | |
| import threading | |
| import uuid | |
| import torch | |
| from flask import Flask, jsonify, request, send_file, redirect | |
| from flask_cors import CORS | |
| os.environ["OMP_NUM_THREADS"] = "4" | |
| os.environ["MKL_NUM_THREADS"] = "4" | |
| torch.set_num_threads(4) | |
| torch.set_num_interop_threads(1) | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s") | |
| logger = logging.getLogger("sd-turbo-server") | |
| DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "4")) | |
| DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "1.0")) | |
| DEFAULT_WIDTH = 768 | |
| DEFAULT_HEIGHT = 768 | |
| app = Flask(__name__) | |
| CORS(app) | |
| _pipeline = None | |
| _pipeline_lock = threading.Lock() | |
| _pipeline_load_error = None | |
| def load_pipeline(): | |
| global _pipeline, _pipeline_load_error | |
| from diffusers import AutoPipelineForText2Image | |
| model_dir = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo") | |
| logger.info("Loading pipeline from %s", model_dir) | |
| try: | |
| pipeline = AutoPipelineForText2Image.from_pretrained( | |
| model_dir, torch_dtype=torch.float32, safety_checker=None | |
| ) | |
| pipeline.to("cpu") | |
| pipeline.unet = torch.quantization.quantize_dynamic( | |
| pipeline.unet, {torch.nn.Linear}, dtype=torch.qint8 | |
| ) | |
| if hasattr(pipeline, "text_encoder"): | |
| pipeline.text_encoder = torch.quantization.quantize_dynamic( | |
| pipeline.text_encoder, {torch.nn.Linear}, dtype=torch.qint8 | |
| ) | |
| pipeline.set_progress_bar_config(disable=True) | |
| _pipeline = pipeline | |
| logger.info("Pipeline loaded and quantized successfully.") | |
| except Exception as e: | |
| _pipeline_load_error = str(e) | |
| logger.error("Pipeline load failed: %s", e) | |
| def run_generation(prompt, width, height, steps, guidance): | |
| with _pipeline_lock: | |
| result = _pipeline( | |
| prompt=prompt, | |
| num_inference_steps=steps, | |
| guidance_scale=guidance, | |
| width=width, | |
| height=height, | |
| ) | |
| return result.images[0] | |
| def index(): | |
| return "SD Turbo Server - Use /image or /v1/images/generations" | |
| def health(): | |
| return jsonify({"status": "ok", "model": "sd-turbo"}) | |
| def list_models(): | |
| return jsonify({ | |
| "object": "list", | |
| "data": [{"id": "sd-turbo", "object": "model", "owned_by": "stabilityai"}] | |
| }) | |
| def images_generations(): | |
| """OpenAI-compatible endpoint that returns a direct image URL""" | |
| if not request.is_json: | |
| return jsonify({"error": {"message": "JSON required"}}), 400 | |
| data = request.get_json() | |
| prompt = data.get("prompt") | |
| if not prompt: | |
| return jsonify({"error": {"message": "prompt is required"}}), 400 | |
| size = data.get("size", "768x768") | |
| n = data.get("n", 1) | |
| steps = int(data.get("num_inference_steps", DEFAULT_STEPS)) | |
| guidance = float(data.get("guidance_scale", DEFAULT_GUIDANCE)) | |
| if steps > 6: steps = 6 | |
| if guidance > 2.0: guidance = 1.5 | |
| width = DEFAULT_WIDTH | |
| height = DEFAULT_HEIGHT | |
| if "x" in size: | |
| try: | |
| w, h = map(int, size.split("x")) | |
| width, height = min(w, 768), min(h, 768) | |
| except: | |
| pass | |
| request_id = uuid.uuid4().hex[:8] | |
| logger.info(f"[{request_id}] OpenAI gen - prompt={prompt[:60]}... size={width}x{height}") | |
| try: | |
| image = run_generation(prompt, width, height, steps, guidance) | |
| # Save to a temporary in-memory route so we can return a URL | |
| img_id = uuid.uuid4().hex[:12] | |
| # For simplicity we return a direct /image URL with the same prompt | |
| image_url = f"https://cloudunity-sdturbolumi.hf.space/image?prompt={prompt.replace(' ', '+')}&width={width}&height={height}&steps={steps}&guidance={guidance}" | |
| return jsonify({ | |
| "created": int(time.time()), | |
| "data": [{"url": image_url}] | |
| }) | |
| except Exception as e: | |
| logger.exception("Generation failed") | |
| return jsonify({"error": {"message": str(e)}}), 500 | |
| def generate_image(): | |
| prompt = request.args.get("prompt", "").replace("+", " ") | |
| if not prompt: | |
| return "Missing prompt", 400 | |
| try: | |
| width = int(request.args.get("width", DEFAULT_WIDTH)) | |
| height = int(request.args.get("height", DEFAULT_HEIGHT)) | |
| steps = int(request.args.get("steps", DEFAULT_STEPS)) | |
| guidance = float(request.args.get("guidance", DEFAULT_GUIDANCE)) | |
| except ValueError: | |
| return "Invalid parameters", 400 | |
| if width > 768: width = 768 | |
| if height > 768: height = 768 | |
| if steps > 6: steps = 6 | |
| if guidance > 2.0: guidance = 1.0 | |
| try: | |
| img = run_generation(prompt, width, height, steps, guidance) | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| buf.seek(0) | |
| return send_file(buf, mimetype="image/png") | |
| except Exception as e: | |
| logger.exception("Image generation failed") | |
| return str(e), 500 | |
| if __name__ == "__main__": | |
| load_pipeline() | |
| port = int(os.environ.get("PORT", "7860")) | |
| app.run(host="0.0.0.0", port=port, threaded=True) |