/worker-sdxl-lightning

Template for building custom RunPod Endpoint API workers using SDXL Lightning for image generation.

Primary LanguagePythonMIT LicenseMIT

SDXL Lightning Worker Template

A specialized worker template for building custom RunPod Endpoint API workers utilizing the SDXL Lightning model.

Setup

Prerequisites: Docker & git installed, an sdxl-lightning model as .safetensors file

  1. test the model pipeline/config in src/handler.py locally and make adjustments if necessary
  2. git clone https://github.com/notune/worker-sdxl-lightning.git
  3. cd worker-sdxl-lightning
  4. Add model.safetensors to root directory (optional: if you named it differently change the line ADD model.safetensors / in Dockerfile and pipe = StableDiffusionXLPipeline.from_single_file("/model.safetensors", torch_dtype=torch.float16, variant="fp16") in src/handler.py)
  5. sudo DOCKER_BUILDKIT=1 docker build .
  6. For Docker Hub: Create Repo on https://hub.docker.com/
  7. docker login --username=yourhubusername
  8. Copy image-id from recently built image: docker images
  9. Tag image: docker tag <image-id> yourhubusername/sdxll-custom:1.0.0
  10. docker push yourhubusername/sdxll-custom:1.0.0
  11. create serverless worker on runpod with container image name yourhubusername/sdxll-custom:1.0.0

Example Request

import requests
import base64
from PIL import Image
import io

url = "https://api.runpod.ai/v2/<YOUR_RUNPOD_ID>/runsync"
headers = {
    "Content-Type": "application/json",
    "Authorization": "Bearer <YOUR_TOKEN_HERE>"
}
data = {
    "input": {
        "prompt": "a cute 3d animated cat, pixar style"
    }
}

response = requests.post(url, headers=headers, json=data)

# Check if the response is successful
if response.status_code == 200:
    response_data = response.json()
    if response_data['status'] == 'COMPLETED':
        # Decode the base64 image data
        base64_string = response_data['output']
        image_data = base64.b64decode(base64_string)
        image = Image.open(io.BytesIO(image_data))
        image.show()  # Display the image
        image.save("output_image.png")  # Save the image as a PNG file
    else:
        print("Request was not completed successfully.")
else:
    print(f"Failed to get a valid response. Status code: {response.status_code}")