|
import io |
|
import os |
|
import random |
|
import string |
|
import sys |
|
import torch |
|
from diffusers import StableDiffusionPipeline, DDIMScheduler |
|
import numpy as np |
|
import requests |
|
from PIL import Image |
|
from io import BytesIO |
|
import pydash as _ |
|
import boto3 |
|
import time |
|
import subprocess |
|
from GroundedSegmentAnything import installation |
|
|
|
S3_REGION = "fra1" |
|
S3_ACCESS_ID = "0RN7BZXS59HYSBD3VB79" |
|
S3_ACCESS_SECRET = "hfSPgBlWl5jsGHa2xuByVkSpancgVeA2CVQf2EMp" |
|
S3_ENDPOINT_URL = "https://s3.solarcom.ch" |
|
S3_BUCKET_NAME = "pissnelke" |
|
|
|
s3_session = boto3.session.Session() |
|
s3 = s3_session.client( |
|
service_name="s3", |
|
region_name=S3_REGION, |
|
aws_access_key_id=S3_ACCESS_ID, |
|
aws_secret_access_key=S3_ACCESS_SECRET, |
|
endpoint_url=S3_ENDPOINT_URL, |
|
) |
|
|
|
get_nude_function = None |
|
|
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
|
|
os.environ['path'] = path |
|
print("running apt-get update && apt-get install ffmpeg libsm6 libxext6 -y") |
|
command = "apt-get update && apt-get install ffmpeg libsm6 libxext6 -y" |
|
process = subprocess.Popen( |
|
command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) |
|
stdout, stderr = process.communicate() |
|
print("ran apt-get update && apt-get install ffmpeg libsm6 libxext6 -y") |
|
print("path", path) |
|
global get_nude_function |
|
from main import get_nude |
|
get_nude_function = get_nude |
|
|
|
def get_pipe(self, loras, lora_weights): |
|
pipe = _.clone_deep(self.base_pipe) |
|
|
|
lora_names = _.map_(loras, lambda x: x.split(".")[0]) |
|
lora_weights = _.map_(lora_weights, lambda x: float(x)) |
|
|
|
if len(lora_weights) > 0 and len(lora_weights) == len(lora_names): |
|
pipe.set_adapters(lora_names, adapter_weights=lora_weights) |
|
|
|
return pipe |
|
|
|
def __call__(self, data): |
|
original_image_res = requests.get(data.get("original_link")) |
|
original_pil = Image.open(BytesIO(original_image_res.content)) |
|
|
|
replicate_api_key = data.get("replicate_api_key", "") |
|
|
|
with_small_tits = data.get("with_small_tits", False) |
|
|
|
with_big_tits = data.get("with_big_tits", False) |
|
|
|
nude_pils = get_nude_function(cfg_scale=data.get("cfg_scale"), generate_max_size=data.get("generate_max_size"), original_max_size=data.get( |
|
"original_max_size"), original_pil=original_pil, positive_prompt=data.get("positive_prompt"), steps=data.get("steps"), replicate_api_key=replicate_api_key, with_small_tits=with_small_tits, with_big_tits=with_big_tits) |
|
|
|
filenames = [] |
|
|
|
for image in nude_pils: |
|
byte_arr = io.BytesIO() |
|
image.save(byte_arr, format='PNG') |
|
byte_arr = byte_arr.getvalue() |
|
|
|
random_string = ''.join(random.choice( |
|
string.ascii_letters + string.digits) for i in range(20)) |
|
image_filename = random_string + ".jpeg" |
|
|
|
s3.put_object(Body=byte_arr, Bucket=S3_BUCKET_NAME, |
|
Key=image_filename) |
|
|
|
filenames.append(image_filename) |
|
|
|
return { |
|
"filenames": filenames |
|
} |
|
|
