Create api.py

api.py

@@ -0,0 +1,252 @@
+import os
+import glob
+import torch
+from flask import Flask, request, jsonify, send_file
+from diffusers import StableDiffusionPipeline, AutoencoderKL
+from PIL import Image, ImageFilter
+from deep_translator import GoogleTranslator
+
+print("IMAGE PRO AI STARTING...")
+
+# ===============================
+# PATHS
+# ===============================
+
+CACHE_DIR = "F:/AI_models"
+OUTPUT_DIR = "outputs"
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# ===============================
+# TRANSLATOR
+# ===============================
+
+def translate_prompt(prompt):
+
+    try:
+
+        prompt_en = GoogleTranslator(
+            source="pl",
+            target="en"
+        ).translate(prompt)
+
+        print("PL:", prompt)
+        print("EN:", prompt_en)
+
+        return prompt_en
+
+    except:
+
+        return prompt
+
+# ===============================
+# FIND VAE
+# ===============================
+
+print("Searching VAE...")
+
+vae = None
+
+vae_files = glob.glob(
+r"F:\AI_models\models--stabilityai--sd-vae-ft-mse\snapshots\*\vae-ft-mse-840000-ema-pruned.safetensors"
+)
+
+if len(vae_files) > 0:
+
+    VAE_PATH = vae_files[0]
+
+    print("VAE FOUND:", VAE_PATH)
+
+    try:
+
+        vae = AutoencoderKL.from_single_file(
+            VAE_PATH,
+            torch_dtype=torch.float32
+        )
+
+    except:
+
+        print("VAE LOAD FAILED")
+
+# ===============================
+# LOAD MODEL
+# ===============================
+
+print("Loading AI model...")
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "SG161222/Realistic_Vision_V5.1_noVAE",
+    cache_dir=CACHE_DIR,
+    torch_dtype=torch.float32,
+    safety_checker=None,
+    vae=vae
+)
+
+pipe.enable_attention_slicing()
+
+print("MODEL READY")
+
+# ===============================
+# FLASK
+# ===============================
+
+app = Flask(__name__)
+
+# ======================================================
+# 1 GENERATE IMAGE
+# ======================================================
+
+@app.route("/generate", methods=["POST"])
+def generate():
+
+    data = request.json
+
+    prompt = translate_prompt(data["prompt"])
+
+    image = pipe(
+        prompt,
+        num_inference_steps=30,
+        guidance_scale=7.5
+    ).images[0]
+
+    path = OUTPUT_DIR + "/gen.png"
+    image.save(path)
+
+    return send_file(path, mimetype="image/png")
+
+# ======================================================
+# 2 PRODUCT AD (2 STYLE DNA)
+# ======================================================
+
+@app.route("/product_ad", methods=["POST"])
+def product_ad():
+
+    data = request.json
+    prompt = translate_prompt(data["prompt"])
+
+    style1 = prompt + ", luxury product photography, studio lighting"
+    style2 = prompt + ", commercial advertising, dramatic lighting"
+
+    img1 = pipe(style1).images[0]
+    img2 = pipe(style2).images[0]
+
+    p1 = OUTPUT_DIR + "/ad1.png"
+    p2 = OUTPUT_DIR + "/ad2.png"
+
+    img1.save(p1)
+    img2.save(p2)
+
+    return jsonify({
+        "image1": p1,
+        "image2": p2
+    })
+
+# ======================================================
+# 3 FIX OLD PHOTO
+# ======================================================
+
+@app.route("/restore", methods=["POST"])
+def restore():
+
+    file = request.files["image"]
+
+    img = Image.open(file.stream)
+
+    img = img.filter(ImageFilter.SHARPEN)
+
+    path = OUTPUT_DIR + "/restored.png"
+
+    img.save(path)
+
+    return send_file(path, mimetype="image/png")
+
+# ======================================================
+# 4 HD UPSCALE
+# ======================================================
+
+@app.route("/upscale", methods=["POST"])
+def upscale():
+
+    file = request.files["image"]
+
+    img = Image.open(file.stream)
+
+    width, height = img.size
+
+    img = img.resize(
+        (width*2, height*2),
+        Image.LANCZOS
+    )
+
+    path = OUTPUT_DIR + "/hd.png"
+
+    img.save(path)
+
+    return send_file(path, mimetype="image/png")
+
+# ======================================================
+# 5 COLORIZE
+# ======================================================
+
+@app.route("/colorize", methods=["POST"])
+def colorize():
+
+    file = request.files["image"]
+
+    img = Image.open(file.stream)
+
+    img = img.convert("RGB")
+
+    path = OUTPUT_DIR + "/color.png"
+
+    img.save(path)
+
+    return send_file(path, mimetype="image/png")
+
+# ======================================================
+# MEGA UPGRADE — STUDIO PRO
+# ======================================================
+
+@app.route("/studio_pro", methods=["POST"])
+def studio_pro():
+
+    data = request.json
+    prompt = translate_prompt(data["prompt"])
+
+    prompts = [
+
+    prompt + ", luxury studio lighting",
+    prompt + ", product on marble table",
+    prompt + ", premium advertising photo",
+    prompt + ", instagram product ad",
+    prompt + ", minimal product photography"
+
+    ]
+
+    results = []
+
+    for p in prompts:
+
+        img = pipe(p).images[0]
+
+        path = OUTPUT_DIR + "/" + str(len(results)) + ".png"
+
+        img.save(path)
+
+        results.append(path)
+
+    return jsonify(results)
+
+# ======================================================
+# START SERVER
+# ======================================================
+
+if __name__ == "__main__":
+
+    print("IMAGE PRO API READY")
+
+    app.run(
+        host="0.0.0.0",
+        port=5000,
+        debug=False
+    )