MesyAI / main.py
thejagstudio's picture
Update main.py
06c5071 verified
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)