Spaces:
Running
Running
File size: 1,343 Bytes
2c50826 34046e2 2c50826 34046e2 2c50826 |
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 |
import base64
import io
import time
from pathlib import Path
from typing import Any
from dotenv import load_dotenv
from PIL import Image
from together import Together
from api.flux import FluxAPI
class TogetherAPI(FluxAPI):
def __init__(self):
load_dotenv()
self._client = Together()
@property
def name(self) -> str:
return "together"
def generate_image(self, prompt: str, save_path: Path) -> float:
start_time = time.time()
result = self._client.images.generate(
prompt=prompt,
model="black-forest-labs/FLUX.1-dev",
width=1024,
height=1024,
steps=28,
n=1,
seed=0,
response_format="b64_json",
)
end_time = time.time()
if result and hasattr(result, "data") and len(result.data) > 0:
self._save_image_from_result(result, save_path)
else:
raise Exception("No result returned from Together API")
return end_time - start_time
def _save_image_from_result(self, result: Any, save_path: Path):
save_path.parent.mkdir(parents=True, exist_ok=True)
b64_str = result.data[0].b64_json
image_data = base64.b64decode(b64_str)
image = Image.open(io.BytesIO(image_data))
image.save(save_path)
|