Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,515 Bytes
831c617 |
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 |
import runpod
import auth
import app_imagegen
import json
import uuid
import boto3 # Ensure you have boto3 installed for Cloudflare R2 interaction
from botocore.exceptions import NoCredentialsError
from botocore.client import Config
import os # Import os to access environment variables
from datetime import datetime
import re
config = json.load(open("config.json"))
model_manager = app_imagegen.ModelManager(config, max_loaded_models=2, lazy_load=True)
def handler(job):
job_input = job["input"]
result = app_imagegen.image_gen_api_entry(
job_input["prompt"],
job_input.get("negative_prompt",""),
job_input.get("styles",{}),
job_input.get("options",{}),
job_input.get("token",""),
model_manager
)
user_id = job_input.get("user_id","")
domain = job_input.get("domain","")
# Save result as WebP using a GUID as filename
filename = f"{uuid.uuid4()}.webp"
with open(filename, "wb") as f:
f.write(result) # Assuming result is in bytes
# Upload to Cloudflare R2
r2_client = boto3.client('s3',
endpoint_url=f"https://{config['r2']['account_id']}.r2.cloudflarestorage.com",
aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'), # Read from environment variable
aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'), # Read from environment variable
config=Config(signature_version='s3v4'),
region_name='auto'
)
# Prepare metadata
publish_time = datetime.now().isoformat()
sanitized_prompt = re.sub(r'[^\x20-\x7E]', ' ', job_input["prompt"]) # Replace non-ASCII characters with space
sanitized_prompt = sanitized_prompt[:512] # Limit length to 256 characters
metadata = {
'domain': domain,
'model': job_input.get("styles",{}).get("vibe","") ,
'prompt': ''.join([i if ord(i) < 128 else ' ' for i in sanitized_prompt]),
'author': user_id,
'publish_time': str(publish_time)
}
new_url_path = ""
try:
r2_key = f"paid/{user_id}/{filename}"
r2_client.put_object(
Bucket=config['r2']['bucket_name'],
Key=r2_key,
Body=open(filename, 'rb'),
ContentType='image/webp', # Adjust this if the image type might vary
Metadata=metadata
)
new_url_path = f"/{r2_key}"
except NoCredentialsError:
print("Credentials not available")
return {"image_url":new_url_path}
runpod.serverless.start({"handler": handler}) |