File size: 3,150 Bytes
72dddd7 0893e31 72dddd7 176d2d4 72dddd7 176d2d4 72dddd7 176d2d4 72dddd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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=""):
# get_nude(Image.open("girl.png"))
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
}
|