Update main.py

main.py

@@ -1,2 +1,386 @@
-import os
-exec(os.environ.get('CODE'))
+from flask import Flask,jsonify,request,send_file, render_template
+from flask_cors import CORS, cross_origin
+import requests
+import json
+from io import BytesIO
+import base64
+
+from transformers import SegformerImageProcessor, TFSegformerForSemanticSegmentation
+from PIL import Image,ImageFilter,ImageEnhance
+import matplotlib.pyplot as plt
+import tensorflow as tf
+import numpy as np
+
+processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+model = TFSegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+
+
+app = Flask(__name__)
+cors = CORS(app)
+app.config['CORS_HEADERS'] = 'Content-Type'
+
+def ade_palette():
+    return [[0,0,0],[0,0,0],[0,0,0],[0,0,0],[255,255,255],[255,255,255],[255,255,255],[255,255,255],[255,255,255],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[255,255,255]]
+
+colormap = np.asarray(ade_palette())
+
+
+@app.route('/api/v1/sdxl/generate', methods=['post'])
+def generator():
+    try:
+        data = request.get_json()
+        cameraAngle = data["cameraAngle"]
+        age = data["age"]
+        base = data["base"]
+        expression = data["expression"]
+        bodyType = data["bodyType"]
+        style = data["style"]
+        action = data["action"]
+        boobs = data["boobs"]
+        ass = data["ass"]
+        clothing = data["clothing"]
+        hairStyle = data["hairStyle"]
+        quality = data["quality"]
+        clothingColor = data["clothingColor"]
+        clothingType = data["clothingType"]
+        clothingMaterial = data["clothingMaterial"]
+        clothingPattern = data["clothingPattern"]
+        lighting = data["lighting"]
+        setting = data["setting"]
+        background = data["background"]
+        seed = int(data["seed"])
+
+        url = "https://api.prodia.com/v1/sdxl/generate"
+        
+        prompt ="A "+cameraAngle+" portrait of a "+age+" "+base+" with a "+expression+" and "+bodyType+" body, "+action+", "+boobs+" boobs, "+ass+" ass,  "+clothing+", "+hairStyle+" hair, "+quality+" quality, "+clothingColor+" "+clothingType+", "+clothingMaterial+" with "+clothingPattern+", "+lighting+" in "+setting+" with "+background+" in background"
+        model = "dreamshaperXL10_alpha2.safetensors [c8afe2ef]"
+        if style == "photographic":
+            model = "juggernautXL_v45.safetensors [e75f5471]"
+        # model = "sd_xl_base_1.0.safetensors [be9edd61]"
+        
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "seed": seed,
+            "cfg_scale": 15,
+        }
+        if style != "":
+            payload["style_preset"] = style
+        payload = json.dumps(payload)
+        headers = {
+            'X-Prodia-Key': 'f2c81250-a5d8-461e-9caa-9fdb5abb5b60',
+            'accept': 'application/json',
+            'content-type': 'application/json'
+        }
+        
+        response = requests.request("POST", url, headers=headers, data=payload)
+        try:
+            job = response.json()
+            return jsonify(job)
+        except Exception as e:
+            return jsonify(error=response.text,source="generator")
+    except Exception as e:
+        return jsonify(error=str(e),source="server")
+
+@app.route('/api/v1/image/<jobId>', methods=['get'])
+def image(jobId):
+    try:
+        url = "https://api.prodia.com/v1/job/"+ jobId.split(".")[0]
+
+        headers = {
+            'X-Prodia-Key': 'f2c81250-a5d8-461e-9caa-9fdb5abb5b60',
+            'accept': 'application/json'
+        }
+
+        responseImg = requests.request("GET", url, headers=headers)
+        if responseImg.json()["status"] == "succeeded":
+
+            url = "https://images.prodia.xyz/"+jobId
+            headers = {
+                'accept': 'image/png'
+            }
+            response = requests.request("GET", url, headers=headers)
+            responseMain = app.response_class(
+                response=response.content,
+                status=200,
+                mimetype='image/png'
+            )
+            return responseMain
+        elif responseImg.json()["status"] == "failed":
+            url = "https://www.elegantthemes.com/blog/wp-content/uploads/2020/08/000-http-error-codes.png"
+            headers = {
+                'accept': 'image/png'
+            }
+            response = requests.request("GET", url, headers=headers)
+            responseMain = app.response_class(
+                response=response.content,
+                status=200,
+                mimetype='image/png'
+            )
+            return responseMain
+        else:
+            return jsonify(responseImg.json())
+    except Exception as e:
+        return jsonify(error=str(e),source="server")
+
+@app.route('/api/v1/nudify', methods=['post'])
+def nudify():
+    global model,processor
+    try:
+        data = request.get_json()
+        if "/" not in data["image"]:
+            url = "https://images.prodia.xyz/"+ data["image"]
+        else:
+            url = data["image"]
+        cameraAngle = data["cameraAngle"]
+        age = data["age"]
+        base = data["base"]
+        expression = data["expression"]
+        bodyType = data["bodyType"]
+        style = data["style"]
+        action = data["action"]
+        boobs = data["boobs"]
+        ass = data["ass"]
+        clothing = data["clothing"]
+        hairStyle = data["hairStyle"]
+        quality = data["quality"]
+        clothingColor = data["clothingColor"]
+        clothingType = data["clothingType"]
+        clothingMaterial = data["clothingMaterial"]
+        clothingPattern = data["clothingPattern"]
+        lighting = data["lighting"]
+        setting = data["setting"]
+        background = data["background"]
+        seed = int(data["seed"])
+        
+        prompt ="A "+cameraAngle+" angle full naked portrait of a "+age+" aged "+base+" with a "+bodyType+" body naked, "+action+", "+boobs+" boobs, "+ass+" ass, nude, naked, nsfw"
+        modelGen = "v1-5-inpainting.safetensors [21c7ab71]"
+        # modelGen = "sd_xl_base_1.0_inpainting_0.1.safetensors [5679a81a]"
+
+        input_img = Image.open(requests.get(url,stream=True).raw)
+        # input_img.save("input.png")
+        inputs = processor(images=input_img,return_tensors="tf")
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+        # upsampled_logits = nn.functional.interpolate(
+        #     logits,
+        #     size=input_img.size[::-1],
+        #     mode="bilinear",
+        #     align_corners=False,
+        # )
+
+        # seg = upsampled_logits.argmax(dim=1)[0]
+        logits = tf.transpose(logits, [0, 2, 3, 1])
+        logits = tf.image.resize(
+            logits, input_img.size[::-1]
+        )  # We reverse the shape of `image` because `image.size` returns width and height.
+        seg = tf.math.argmax(logits, axis=-1)[0]
+
+        color_seg = np.zeros(
+            (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
+        )  # height, width, 3
+        for label, color in enumerate(colormap):
+            color_seg[seg.numpy() == label, :] = color
+
+        # Show image + mask
+        pred_img = color_seg * 1
+        pred_img = pred_img.astype(np.uint8)
+
+        imgOut = Image.fromarray(pred_img)
+        # blur the mask
+        imgOut = imgOut.filter(ImageFilter.GaussianBlur(16))
+        imgOut = imgOut.convert("L")
+        # Apply thresholding to sharpen the image
+        threshold = 20
+        imgOut = imgOut.point(lambda p: p > threshold and 255)
+
+        # Enhance the sharpness of the image
+        enhancer = ImageEnhance.Sharpness(imgOut)
+        imgOut = enhancer.enhance(2.0)
+        imgOut = imgOut.filter(ImageFilter.GaussianBlur(8))
+        imgOut = imgOut.convert("RGB")
+        
+        # imgOut.save("mask.png")
+        buffered = BytesIO()
+        imgOut.save(buffered, format="png")
+        base64_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        url2 = "https://api.prodia.com/v1/sd/inpainting"
+
+        payload = json.dumps({
+            "inpainting_full_res": False,
+            "imageUrl": url,
+            "maskData": base64_str,
+            "model": modelGen,
+            "prompt": prompt,
+            "negative_prompt": "bad anatomy, bad hands, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, worst face, three crus, extra crus, fused crus, worst feet, three feet, fused feet, fused thigh, three thigh, fused thigh, extra thigh, worst thigh, missing fingers, extra fingers, ugly fingers, long fingers, horn, extra eyes, huge eyes, 2girl, amputation, disconnected limbs",
+            "denoising_strength": 0.9,
+            "cfg_scale": 7,
+            "mask_blur": 4,
+            "inpainting_fill": 2,
+            "inpainting_mask_invert": 0,
+            "seed": seed,
+            "style_preset" : style,
+            "upscale": True
+        })
+        headers = {
+            'X-Prodia-Key': 'f2c81250-a5d8-461e-9caa-9fdb5abb5b60',
+            'accept': 'application/json',
+            'content-type': 'application/json'
+        }
+
+        response = requests.request("POST", url2, headers=headers, data=payload)
+
+        try:
+            job = response.json()
+            return jsonify(job)
+        except Exception as e:
+            return jsonify(error=response.text,source="generator")
+    except Exception as e:
+        return jsonify(error=str(e),source="server")
+
+@app.route('/api/v1/clotheChanger', methods=['post'])
+def clotheChanger():
+    global model,processor
+    try:
+        data = request.get_json()
+        if "/" not in data["image"]:
+            url = "https://images.prodia.xyz/"+ data["image"]
+        else:
+            url = data["image"]
+        cameraAngle = data["cameraAngle"]
+        age = data["age"]
+        base = data["base"]
+        expression = data["expression"]
+        bodyType = data["bodyType"]
+        style = data["style"]
+        action = data["action"]
+        boobs = data["boobs"]
+        ass = data["ass"]
+        clothing = data["clothing"]
+        hairStyle = data["hairStyle"]
+        quality = data["quality"]
+        clothingColor = data["clothingColor"]
+        clothingType = data["clothingType"]
+        clothingMaterial = data["clothingMaterial"]
+        clothingPattern = data["clothingPattern"]
+        lighting = data["lighting"]
+        setting = data["setting"]
+        background = data["background"]
+        seed = int(data["seed"])
+        
+        prompt ="A "+cameraAngle+" portrait of a "+age+" "+base+" with a "+expression+" and "+bodyType+" body, "+action+", "+boobs+" boobs, "+ass+" ass,  "+clothing+", "+hairStyle+" hair, "+quality+" quality, "+clothingColor+" "+clothingType+", "+clothingMaterial+" with "+clothingPattern+", "+lighting+" in "+setting+" with "+background+" in background"
+        modelGen = "sd_xl_base_1.0_inpainting_0.1.safetensors [5679a81a]"
+
+        input_img = Image.open(requests.get(url,stream=True).raw)
+        inputs = processor(images=input_img,return_tensors="tf")
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+        logits = tf.transpose(logits, [0, 2, 3, 1])
+        logits = tf.image.resize(
+            logits, input_img.size[::-1]
+        )  # We reverse the shape of `image` because `image.size` returns width and height.
+        seg = tf.math.argmax(logits, axis=-1)[0]
+
+        color_seg = np.zeros(
+            (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
+        )  # height, width, 3
+        for label, color in enumerate(colormap):
+            color_seg[seg.numpy() == label, :] = color
+
+        # Show image + mask
+        pred_img = color_seg * 1
+        pred_img = pred_img.astype(np.uint8)
+
+        imgOut = Image.fromarray(pred_img)
+        # blur the mask
+        imgOut = imgOut.filter(ImageFilter.GaussianBlur(16))
+        imgOut = imgOut.convert("L")
+        # Apply thresholding to sharpen the image
+        threshold = 20
+        imgOut = imgOut.point(lambda p: p > threshold and 255)
+
+        # Enhance the sharpness of the image
+        enhancer = ImageEnhance.Sharpness(imgOut)
+        imgOut = enhancer.enhance(2.0)
+        imgOut = imgOut.filter(ImageFilter.GaussianBlur(8))
+        imgOut = imgOut.convert("RGB")
+        
+        imgOut.save("mask.png")
+        buffered = BytesIO()
+        imgOut.save(buffered, format="png")
+        base64_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        url2 = "https://api.prodia.com/v1/sdxl/inpainting"
+
+        payload = json.dumps({
+            "inpainting_full_res": False,
+            "imageUrl": url,
+            "maskData": base64_str,
+            "model": modelGen,
+            "prompt": prompt,
+            "negative_prompt": "bad anatomy, bad hands, three hands, three legs, bad arms, missing legs, missing arms, poorly drawn face, bad face, fused face, cloned face, worst face, three crus, extra crus, fused crus, worst feet, three feet, fused feet, fused thigh, three thigh, fused thigh, extra thigh, worst thigh, missing fingers, extra fingers, ugly fingers, long fingers, horn, extra eyes, huge eyes, 2girl, amputation, disconnected limbs",
+            "denoising_strength": 0.9,
+            "cfg_scale": 7,
+            "mask_blur": 4,
+            "inpainting_fill": 2,
+            "inpainting_mask_invert": 0,
+            "seed": seed,
+            "style_preset" : style
+        })
+        headers = {
+            'X-Prodia-Key': 'f2c81250-a5d8-461e-9caa-9fdb5abb5b60',
+            'accept': 'application/json',
+            'content-type': 'application/json'
+        }
+
+        response = requests.request("POST", url2, headers=headers, data=payload)
+
+        try:
+            job = response.json()
+            return jsonify(job)
+        except Exception as e:
+            return jsonify(error=response.text,source="generator")
+    except Exception as e:
+        return jsonify(error=str(e),source="server")
+
+@app.route('/api/v1/upscale', methods=['post'])
+def upscale():
+    try:
+        data = request.get_json()
+        if "/" not in data["image"]:
+            imgUrl = "https://images.prodia.xyz/"+ data["image"]
+        else:
+            imgUrl = data["image"]
+        url = "https://api.prodia.com/v1/upscale"
+        headers = {
+                'X-Prodia-Key': 'f2c81250-a5d8-461e-9caa-9fdb5abb5b60',
+                'accept': 'application/json',
+                'content-type': 'application/json',
+            }
+        payload = json.dumps({
+            "resize": 4,
+            "model": "ESRGAN_4x",
+            "imageUrl": imgUrl
+        })
+        response = requests.request("POST", url, headers=headers, data=payload)
+        try:
+            job = response.json()
+            return jsonify(job)
+        except Exception as e:
+            return jsonify(error=response.text,source="generator")
+    except Exception as e:
+        return jsonify(error=str(e),source="server")
+
+@app.route('/', methods=['get'])
+def home():
+    return render_template("index.html")
+@app.route('/generator', methods=['get'])
+def generatorPage():
+    return render_template("index.html")
+
+
+app.run(debug=True, host="0.0.0.0",port=7860)
+# app.run(debug=True, host="127.0.0.1",port=8000)