Spaces:
Running
Running
File size: 1,377 Bytes
2c50826 4f41410 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 49 |
import time
from io import BytesIO
from pathlib import Path
from typing import Any
import fal_client
import requests
from PIL import Image
from api.flux import FluxAPI
class FalAPI(FluxAPI):
@property
def name(self) -> str:
return "fal"
def generate_image(self, prompt: str, save_path: Path) -> float:
start_time = time.time()
result = fal_client.subscribe(
"fal-ai/flux/dev",
arguments={
"seed": 0,
"prompt": prompt,
"image_size": "square_hd", # 1024x1024 image
"num_images": 1,
"guidance_scale": 3.5,
"num_inference_steps": 28,
"enable_safety_checker": True,
},
)
end_time = time.time()
url = result["images"][0]["url"]
self._save_image_from_url(url, save_path)
return end_time - start_time
def _save_image_from_url(self, url: str, save_path: Path):
response = requests.get(url)
image = Image.open(BytesIO(response.content))
save_path.parent.mkdir(parents=True, exist_ok=True)
image.save(save_path)
def _save_image_from_result(self, result: Any, save_path: Path):
save_path.parent.mkdir(parents=True, exist_ok=True)
with open(save_path, "wb") as f:
f.write(result.content)
|