add old files requirements.txt

src/image_gen.py

@@ -1,9 +1,8 @@
-from openai import OpenAI
+import replicate
 from pathlib import Path
 import uuid
 import requests
 from enum import Enum
-from typing import Literal
 
 class Styles(Enum):
     emoji = "emoji"
@@ -11,43 +10,61 @@ class Styles(Enum):
     svg = "svg"
     cartoon = "cartoon"
 
-client = OpenAI()
-
 
 def image_style(base_prompt: str, style: Styles) -> str:
     return f"create {style.name} style image on a white background: {base_prompt}"
 
-def icon_gen(prompt: str, quality: Literal["hd", "standard"], model: Literal["dall-e-3", "dall-e-2"], output: Path = "./output/"):
-    # Make sure output directory exists
-    output_path = Path(output) / model / quality
-    output_path.mkdir(parents=True, exist_ok=True)
 
-    # Generate the image using the OpenAI client
-    response = client.images.generate(
-        model=model,
-        prompt=prompt,
-        size="1024x1024",
-        quality=quality,
-        n=1,
+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])
 
-    # Extract the image URL
-    image_url = response.data[0].url
 
-    # Generate a UUID for the filename
-    image_id = str(uuid.uuid4())
-    image_filename = f"{image_id}.png"
-    image_filepath = output_path / image_filename
+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)
-    print(image_id)
+
+
+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)
-    icon_gen(dog_emoji, model="dall-e-3", output=Path("../output/test"), quality="hd")
+#    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)