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")
# 022ee431-4073-424c-8298-68cb75352785
# e8346eb2-6187-4748-a42f-7241580ee1f1
# ae337a18-427c-4035-a73b-a89b97a14b6c
# c177417d-8b42-423a-8dc1-4b02c7a686ae
# 2347e0bf-ab79-4654-817e-684f0d9e6dff
# d01a7eb8-9f10-4abe-b142-fd5248fb7e88
# 5e6a594a-4bc6-4b06-bed6-5ad57804e307
# 7f6d4a3d-f3e3-4abf-ab8d-a21bb64a5ee7
# 39110d90-70eb-4fed-bb40-3d0bf52c73bf
# 4b35363f-a524-418b-8b2c-ee141e642a8b
# dc80a8a4-0b98-4d54-b3e4-b7c797bc2527
# 5113d4e5-6574-4115-bdd0-e02b5d7cfb87
# aba3968d-b924-45d6-a98f-90ff6b12989c
# da6053eb-c352-4374-a459-2a9a5eaaa64b
# d46c04cb-6ce1-4657-bbfd-8c0c77d6a162
# 06b2ac41-5dc0-461e-85d8-abbac08c5491
# 1008a105-be30-4564-b5d2-edbd4e6a5c3d
# 596e85a1-b3ae-49fb-ae11-cee1c4f78eee
# 3d1ca348-aed2-4fa3-86dd-7583f2f5ef4a
# eaca0864-70a4-4653-8dc7-f5ba3918326f
# fcae63ee-262a-44d9-8941-cdf36ae371ad
# 9f4b2ef5-e6b5-42a9-b459-8958402dfcc5
# 4d012f3f-13f7-4a92-b1f6-ac0e65002803
# 72b18186-fb64-46a4-b4f2-a91af84dd398
# 501eba46-a956-4649-96aa-2d9cc0f048bf
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': 'd2c2605e-b4a5-4ab6-aa32-fcf81e6d27a2',
            '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': 'd2c2605e-b4a5-4ab6-aa32-fcf81e6d27a2',
            '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/sdxl/inpainting"

        payload = json.dumps({
            "inpainting_full_res": False,
            "imageUrl": url,
            "maskData": base64_str,
            "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": 1,
            "cfg_scale": 7,
            "inpainting_fill": 2,
            "inpainting_mask_invert": 0,
            "seed": seed,
            "style_preset" : style,
            "upscale": True,
            "mask_blur": 0,
            "steps": 30
        })
        headers = {
            'X-Prodia-Key': 'd2c2605e-b4a5-4ab6-aa32-fcf81e6d27a2',
            '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 = "v1-5-inpainting.safetensors [21c7ab71]"

        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,
            "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": 1,
            "cfg_scale": 7,
            "inpainting_fill": 2,
            "inpainting_mask_invert": 0,
            "seed": seed,
            "style_preset" : style,
            "upscale": True,
            "mask_blur": 0,
            "steps": 30
        })
        headers = {
            'X-Prodia-Key': 'd2c2605e-b4a5-4ab6-aa32-fcf81e6d27a2',
            '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': 'd2c2605e-b4a5-4ab6-aa32-fcf81e6d27a2',
                '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)