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 70 71 72 73 74 75 76 77 78 79
| import base64 from openai import AsyncOpenAI from pathlib import Path import tempfile import aiohttp import asyncio
class WorkerImage(): def __init__(self, account_id, api_token): self.flux_system_prompt = "...提示词前文有提到,这里不重复写了" self.account_id = account_id self.api_token = api_token self.endpoint = ( f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/@cf/black-forest-labs/flux-1-schnell" ) self.headers = { "Authorization": f"Bearer {api_token}", "Content-Type": "application/json", }
async def _get_flux_image_data(self, prompt, num_steps): payload = { "prompt": prompt, "num_steps": num_steps } async with aiohttp.ClientSession() as session: async with session.post(self.endpoint, headers=self.headers, json=payload) as response: if response.status == 400: return None else: result = await response.json() image_data = base64.b64decode(result['result']['image']) return image_data
def _save_image_to_tmp_file(self, image_data): dir = Path(__file__).resolve().parent with tempfile.NamedTemporaryFile(dir=dir, suffix='.png', delete=False) as f: f.write(image_data) temp_file_path = f.name return temp_file_path
async def get_translate_promt(self, prompt):
client = AsyncOpenAI( api_key=self.api_token, base_url=f"https://api.cloudflare.com/client/v4/accounts/{self.account_id}/ai/v1", max_retries=5, ) chat_completion = await client.chat.completions.create( messages=[ { "role": "system", "content": self.flux_system_prompt, }, { "role": "user", "content": prompt, } ], model="@cf/meta/llama-3.3-70b-instruct-fp8-fast", ) return chat_completion.choices[0].message.content
async def get_image(self, prompt, num_steps, is_translate=False): new_prompts = "" if is_translate: new_prompts = await self.get_translate_promt(prompt) else: new_prompts = prompt image_data = await self._get_flux_image_data(new_prompts, num_steps) if image_data: local_file = self._save_image_to_tmp_file(image_data) local_file = local_file.split("/")[-1] return local_file, new_prompts else: print("image data is None,please check your inputs") return None, new_prompts
|