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 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", "") print("got repli", 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 }