docker-compose build

docker-compose up

from typing import List, Union, Generator, Iterator
from pydantic import BaseModel
import os
import requests

class Pipeline:
 class Valves(BaseModel):
     API_BASE_URL: str = "https://api.cloudflare.com/client/v4/accounts/your_account_id/ai/v1"
     API_KEY: str = ""

 def __init__(self):
     self.type = "manifold"
     self.name = "WorkersAI: "

     self.valves = self.Valves(
         **{
             "API_KEY": os.getenv("API_KEY", "your-cloudflare-api-key-here")
         }
     )
     self.pipelines = self.get_models()

 async def on_startup(self):
     print(f"on_startup:{__name__}")

 async def on_shutdown(self):
     print(f"on_shutdown:{__name__}")

 async def on_valves_updated(self):
     print(f"on_valves_updated:{__name__}")
     self.pipelines = self.get_models()

 def get_models(self):
     return [
         {
             "id": "@cf/bytedance/stable-diffusion-xl-lightning",
             "name": "Stable Diffusion XL Lightning"
         }
     ]

 def pipe(
     self, user_message: str, model_id: str, messages: List[dict], body: dict
 ) -> Union[str, Generator, Iterator]:
     print(f"pipe:{__name__}")

     headers = {
         "Authorization": f"Bearer {self.valves.API_KEY}",
         "Content-Type": "application/json"
     }

     payload = {
         "model": model_id,
         "inputs": {
             "prompt": user_message
         }
     }

     try:
         response = requests.post(
             url=f"{self.valves.API_BASE_URL}/run",
             json=payload,
             headers=headers
         )
         response.raise_for_status()

         if response.headers.get('Content-Type') == 'image/png':
             image_url = response.json().get("url")
             yield f"![image]({image_url})\n"
         else:
             yield response.text

     except Exception as e:
         return f"Error: {e}"