Update api.py

api.py

@@ -1,18 +1,13 @@
 import os
-import glob
 import torch
-from flask import Flask, request, jsonify, send_file, render_template_string
-from diffusers import StableDiffusionPipeline, AutoencoderKL
+from flask import Flask, request, send_file, jsonify, render_template_string
+from diffusers import StableDiffusionPipeline
 from PIL import Image, ImageFilter
 from deep_translator import GoogleTranslator
 
 print("IMAGE PRO AI STARTING...")
 
-# ===============================
-# PATHS
-# ===============================
-
-CACHE_DIR = "F:/AI_models"
+CACHE_DIR = "./models"
 OUTPUT_DIR = "outputs"
 
 os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -22,6 +17,7 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 # ===============================
 
 def translate_prompt(prompt):
+
     try:
         prompt_en = GoogleTranslator(source="pl", target="en").translate(prompt)
         print("PL:", prompt)
@@ -30,44 +26,17 @@ def translate_prompt(prompt):
     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...")
+print("Loading model...")
 
 pipe = StableDiffusionPipeline.from_pretrained(
     "SG161222/Realistic_Vision_V5.1_noVAE",
     cache_dir=CACHE_DIR,
     torch_dtype=torch.float32,
-    safety_checker=None,
-    vae=vae
+    safety_checker=None
 )
 
 pipe.enable_attention_slicing()
@@ -80,76 +49,110 @@ print("MODEL READY")
 
 app = Flask(__name__)
 
-# ======================================================
-# WEB PANEL (TEST GENERATORA)
-# ======================================================
+# ===============================
+# HTML PANEL
+# ===============================
+
+HTML = """
 
-HTML_PAGE = """
-<!DOCTYPE html>
 <html>
+
 <head>
+
 <title>IMAGE PRO AI</title>
+
 <style>
-body{font-family:Arial;background:#111;color:white;text-align:center;padding:40px}
-input{width:400px;padding:10px}
-button{padding:10px 20px;background:#ff6a00;border:none;color:white;font-size:16px}
-img{margin-top:20px;max-width:512px}
+
+body{
+background:black;
+color:white;
+font-family:Arial;
+text-align:center;
+}
+
+input{
+width:400px;
+padding:10px;
+}
+
+button{
+padding:10px 20px;
+background:orange;
+border:none;
+color:white;
+font-size:16px;
+cursor:pointer;
+}
+
+img{
+margin-top:20px;
+max-width:512px;
+}
+
 </style>
+
 </head>
+
 <body>
 
 <h1>IMAGE PRO AI</h1>
 
-<form action="/generate_web" method="post">
-<input name="prompt" placeholder="Wpisz prompt (np. kot w butach)">
+<input id="prompt" placeholder="np. truskawka na stole"/>
 <br><br>
-<button type="submit">GENERUJ</button>
-</form>
 
-{% if image %}
+<button onclick="generate()">GENERUJ</button>
+
 <h2>Wynik</h2>
-<img src="{{image}}">
-{% endif %}
 
-</body>
-</html>
-"""
+<img id="result"/>
 
-@app.route("/", methods=["GET"])
-def home():
-    return render_template_string(HTML_PAGE)
+<script>
 
-# ======================================================
-# GENERATE FROM WEB PANEL
-# ======================================================
+async function generate(){
 
-@app.route("/generate_web", methods=["POST"])
-def generate_web():
+prompt = document.getElementById("prompt").value
 
-    prompt = request.form["prompt"]
-    prompt = translate_prompt(prompt)
+response = await fetch("/generate",{
 
-    image = pipe(
-        prompt,
-        num_inference_steps=30,
-        guidance_scale=7.5
-    ).images[0]
+method:"POST",
 
-    path = OUTPUT_DIR + "/web_gen.png"
-    image.save(path)
+headers:{
+"Content-Type":"application/json"
+},
+
+body:JSON.stringify({
+prompt:prompt
+})
+
+})
+
+blob = await response.blob()
+
+url = URL.createObjectURL(blob)
 
-    return render_template_string(
-        HTML_PAGE,
-        image="/image/web_gen.png"
-    )
+document.getElementById("result").src = url
 
-@app.route("/image/<filename>")
-def serve_image(filename):
-    return send_file(OUTPUT_DIR + "/" + filename, mimetype="image/png")
+}
 
-# ======================================================
-# API GENERATE
-# ======================================================
+</script>
+
+</body>
+
+</html>
+
+"""
+
+# ===============================
+# HOME
+# ===============================
+
+@app.route("/")
+def home():
+    return render_template_string(HTML)
+
+# ===============================
+# GENERATE IMAGE
+# ===============================
 
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -164,40 +167,35 @@ def generate():
     ).images[0]
 
     path = OUTPUT_DIR + "/gen.png"
+
     image.save(path)
 
     return send_file(path, mimetype="image/png")
 
-# ======================================================
+# ===============================
 # PRODUCT AD
-# ======================================================
+# ===============================
 
 @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"
+    style = prompt + ", luxury product photography"
 
-    img1 = pipe(style1).images[0]
-    img2 = pipe(style2).images[0]
+    image = pipe(style).images[0]
 
-    p1 = OUTPUT_DIR + "/ad1.png"
-    p2 = OUTPUT_DIR + "/ad2.png"
+    path = OUTPUT_DIR + "/product.png"
 
-    img1.save(p1)
-    img2.save(p2)
+    image.save(path)
 
-    return jsonify({
-        "image1": p1,
-        "image2": p2
-    })
+    return send_file(path, mimetype="image/png")
 
-# ======================================================
+# ===============================
 # RESTORE PHOTO
-# ======================================================
+# ===============================
 
 @app.route("/restore", methods=["POST"])
 def restore():
@@ -205,6 +203,7 @@ def restore():
     file = request.files["image"]
 
     img = Image.open(file.stream)
+
     img = img.filter(ImageFilter.SHARPEN)
 
     path = OUTPUT_DIR + "/restored.png"
@@ -213,9 +212,9 @@ def restore():
 
     return send_file(path, mimetype="image/png")
 
-# ======================================================
-# UPSCALE HD
-# ======================================================
+# ===============================
+# UPSCALE
+# ===============================
 
 @app.route("/upscale", methods=["POST"])
 def upscale():
@@ -224,12 +223,9 @@ def upscale():
 
     img = Image.open(file.stream)
 
-    width, height = img.size
+    w,h = img.size
 
-    img = img.resize(
-        (width*2, height*2),
-        Image.LANCZOS
-    )
+    img = img.resize((w*2,h*2), Image.LANCZOS)
 
     path = OUTPUT_DIR + "/hd.png"
 
@@ -237,9 +233,9 @@ def upscale():
 
     return send_file(path, mimetype="image/png")
 
-# ======================================================
+# ===============================
 # COLORIZE
-# ======================================================
+# ===============================
 
 @app.route("/colorize", methods=["POST"])
 def colorize():
@@ -247,6 +243,7 @@ def colorize():
     file = request.files["image"]
 
     img = Image.open(file.stream)
+
     img = img.convert("RGB")
 
     path = OUTPUT_DIR + "/color.png"
@@ -255,47 +252,15 @@ def colorize():
 
     return send_file(path, mimetype="image/png")
 
-# ======================================================
-# 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")
+    print("SERVER READY")
 
     app.run(
         host="0.0.0.0",
-        port=7860,
-        debug=False
+        port=7860
     )