70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
import replicate
|
|
from pathlib import Path
|
|
import uuid
|
|
import requests
|
|
from enum import Enum
|
|
|
|
class Styles(Enum):
|
|
emoji = "emoji"
|
|
pixel_art = "pixel art"
|
|
svg = "svg"
|
|
cartoon = "cartoon"
|
|
|
|
|
|
def image_style(base_prompt: str, style: Styles) -> str:
|
|
return f"create {style.name} style image on a white background: {base_prompt}"
|
|
|
|
|
|
def replicate_flux_schnell_image_gen(prompt: str) -> str:
|
|
output = replicate.run(
|
|
"black-forest-labs/flux-schnell",
|
|
input={
|
|
"prompt": prompt,
|
|
"go_fast": True,
|
|
"megapixels": "1",
|
|
"num_outputs": 1,
|
|
"aspect_ratio": "1:1",
|
|
"output_format": "png",
|
|
"output_quality": 80,
|
|
"num_inference_steps": 4
|
|
}
|
|
)
|
|
return str(output[0])
|
|
|
|
|
|
def save_image_url(image_url: str, output_dir: Path):
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
image_filename = str(uuid.uuid4()) + ".png"
|
|
print(image_filename)
|
|
image_filepath = output_dir / image_filename
|
|
|
|
# Download the image
|
|
image_response = requests.get(image_url)
|
|
image_response.raise_for_status()
|
|
with open(image_filepath, "wb") as f:
|
|
f.write(image_response.content)
|
|
|
|
|
|
def replicate_flux_1_1_pro_image_gen(prompt: str) -> str:
|
|
output = replicate.run(
|
|
"black-forest-labs/flux-1.1-pro",
|
|
input={
|
|
"prompt": prompt,
|
|
"aspect_ratio": "1:1",
|
|
"output_format": "png",
|
|
"output_quality": 100,
|
|
"safety_tolerance": 1,
|
|
"prompt_upsampling": True
|
|
}
|
|
)
|
|
return str(output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
base_prompt = "A playful, cartoon-style dog icon featuring a round, fluffy body in a warm golden-brown color. The dog has oversized, expressive black eyes and a happy, open mouth showing its tongue out. Two large floppy ears hang down on either side of its head, each with a slightly darker brown color on the inside. The tail is wagging and curved upwards, represented by a simple, thick shape in the same golden-brown. The background is a light blue with a subtle oval shape to make the dog stand out boldly. This design uses strong shapes and clear contrasts for easy recognition in a small format."
|
|
dog_emoji = image_style(base_prompt, style=Styles.cartoon)
|
|
# image_url = replicate_flux_schnell_image_gen(dog_emoji)
|
|
# output_dir = Path("../output/test2/replicate/flux_schnell")
|
|
image_url = replicate_flux_1_1_pro_image_gen(dog_emoji)
|
|
output_dir = Path("../output/test2/replicate/flux_1_1_pro")
|
|
save_image_url(image_url, output_dir) |