serverless

requirements.txt

@@ -11,4 +11,6 @@ torchsde
 compel
 onnxruntime-gpu
 hidiffusion
+runpod
+boto3
 git+https://github.com/tencent-ailab/IP-Adapter.git

serverless_handler.py

@@ -0,0 +1,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})