refactor: improve code formatting and organization across multiple API and benchmark files

api/__init__.py

@@ -1,31 +1,32 @@
 from typing import Type
 
+from api.aws import AWSBedrockAPI
 from api.baseline import BaselineAPI
+from api.fal import FalAPI
 from api.fireworks import FireworksAPI
 from api.flux import FluxAPI
 from api.pruna import PrunaAPI
 from api.pruna_dev import PrunaDevAPI
 from api.replicate import ReplicateAPI
 from api.together import TogetherAPI
-from api.fal import FalAPI
-from api.aws import AWSBedrockAPI
 
 __all__ = [
-    'create_api',
-    'FluxAPI',
-    'BaselineAPI',
-    'FireworksAPI',
-    'PrunaAPI',
-    'ReplicateAPI',
-    'TogetherAPI',
-    'FalAPI',
-    'PrunaDevAPI',
+    "create_api",
+    "FluxAPI",
+    "BaselineAPI",
+    "FireworksAPI",
+    "PrunaAPI",
+    "ReplicateAPI",
+    "TogetherAPI",
+    "FalAPI",
+    "PrunaDevAPI",
 ]
 
+
 def create_api(api_type: str) -> FluxAPI:
     """
     Factory function to create API instances.
-    
+
     Args:
         api_type (str): The type of API to create. Must be one of:
             - "baseline"
@@ -35,10 +36,10 @@ def create_api(api_type: str) -> FluxAPI:
             - "together"
             - "fal"
             - "aws"
-            
+
     Returns:
         FluxAPI: An instance of the requested API implementation
-        
+
     Raises:
         ValueError: If an invalid API type is provided
     """
@@ -47,7 +48,7 @@ def create_api(api_type: str) -> FluxAPI:
     if api_type.startswith("pruna_"):
         speed_mode = api_type[6:]  # Remove "pruna_" prefix
         return PrunaAPI(speed_mode)
-        
+
     api_map: dict[str, Type[FluxAPI]] = {
         "baseline": BaselineAPI,
         "fireworks": FireworksAPI,
@@ -56,8 +57,10 @@ def create_api(api_type: str) -> FluxAPI:
         "fal": FalAPI,
         "aws": AWSBedrockAPI,
     }
-    
+
     if api_type not in api_map:
-        raise ValueError(f"Invalid API type: {api_type}. Must be one of {list(api_map.keys())} or start with 'pruna_'")
-    
+        raise ValueError(
+            f"Invalid API type: {api_type}. Must be one of {list(api_map.keys())} or start with 'pruna_'"
+        )
+
     return api_map[api_type]()

api/aws.py

@@ -1,9 +1,8 @@
-import os
-import time
 import base64
 import json
+import os
+import time
 from pathlib import Path
-from typing import Any
 
 import boto3
 from dotenv import load_dotenv
@@ -45,23 +44,20 @@ class AWSBedrockAPI(FluxAPI):
         try:
             # Convert request to JSON and invoke the model
             request = json.dumps(native_request)
-            response = self._client.invoke_model(
-                modelId=self._model_id,
-                body=request
-            )
-            
+            response = self._client.invoke_model(modelId=self._model_id, body=request)
+
             # Process the response
             model_response = json.loads(response["body"].read())
             if not model_response.get("images"):
                 raise Exception("No images returned from AWS Bedrock API")
-                
+
             # Save the image
             base64_image_data = model_response["images"][0]
             self._save_image_from_base64(base64_image_data, save_path)
-            
+
         except Exception as e:
             raise Exception(f"Error generating image with AWS Bedrock: {str(e)}")
-            
+
         end_time = time.time()
         return end_time - start_time
 
@@ -70,4 +66,4 @@ class AWSBedrockAPI(FluxAPI):
         save_path.parent.mkdir(parents=True, exist_ok=True)
         image_data = base64.b64decode(base64_data)
         with open(save_path, "wb") as f:
-            f.write(image_data) 
+            f.write(image_data)

api/baseline.py

@@ -12,6 +12,7 @@ class BaselineAPI(FluxAPI):
     """
     As our baseline, we use the Replicate API with go_fast=False.
     """
+
     def __init__(self):
         load_dotenv()
         self._api_key = os.getenv("REPLICATE_API_TOKEN")
@@ -24,6 +25,7 @@ class BaselineAPI(FluxAPI):
 
     def generate_image(self, prompt: str, save_path: Path) -> float:
         import replicate
+
         start_time = time.time()
         result = replicate.run(
             "black-forest-labs/flux-dev",
@@ -39,15 +41,15 @@ class BaselineAPI(FluxAPI):
             },
         )
         end_time = time.time()
-        
+
         if result and len(result) > 0:
             self._save_image_from_result(result[0], save_path)
         else:
             raise Exception("No result returned from Replicate 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)
         with open(save_path, "wb") as f:
-            f.write(result.read()) 
+            f.write(result.read())

api/fal.py

@@ -30,10 +30,10 @@ class FalAPI(FluxAPI):
             },
         )
         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):

api/fireworks.py

@@ -23,7 +23,7 @@ class FireworksAPI(FluxAPI):
 
     def generate_image(self, prompt: str, save_path: Path) -> float:
         start_time = time.time()
-        
+
         headers = {
             "Content-Type": "application/json",
             "Accept": "image/jpeg",
@@ -39,12 +39,12 @@ class FireworksAPI(FluxAPI):
         result = requests.post(self._url, headers=headers, json=data)
 
         end_time = time.time()
-        
+
         if result.status_code == 200:
             self._save_image_from_result(result, save_path)
         else:
             raise Exception(f"Error: {result.status_code} {result.text}")
-        
+
         return end_time - start_time
 
     def _save_image_from_result(self, result: Any, save_path: Path):

api/flux.py

@@ -14,7 +14,7 @@ class FluxAPI(ABC):
     def name(self) -> str:
         """
         The name of the API implementation.
-        
+
         Returns:
             str: The name of the specific API implementation
         """
@@ -24,11 +24,11 @@ class FluxAPI(ABC):
     def generate_image(self, prompt: str, save_path: Path) -> float:
         """
         Generate an image based on the prompt and save it to the specified path.
-        
+
         Args:
             prompt (str): The text prompt to generate the image from
             save_path (Path): The path where the generated image should be saved
-            
+
         Returns:
             float: The time taken for the API call in seconds
         """

api/pruna.py

@@ -3,8 +3,8 @@ import time
 from pathlib import Path
 from typing import Any
 
-from dotenv import load_dotenv
 import replicate
+from dotenv import load_dotenv
 
 from api.flux import FluxAPI
 
@@ -12,7 +12,9 @@ from api.flux import FluxAPI
 class PrunaAPI(FluxAPI):
     def __init__(self, speed_mode: str):
         self._speed_mode = speed_mode
-        self._speed_mode_name = speed_mode.split(" ")[0].strip().lower().replace(" ", "_")
+        self._speed_mode_name = (
+            speed_mode.split(" ")[0].strip().lower().replace(" ", "_")
+        )
         load_dotenv()
         self._api_key = os.getenv("REPLICATE_API_TOKEN")
         if not self._api_key:
@@ -38,7 +40,7 @@ class PrunaAPI(FluxAPI):
             },
         )
         end_time = time.time()
-        
+
         if result:
             self._save_image_from_result(result, save_path)
         else:

api/pruna_dev.py

@@ -3,8 +3,8 @@ import time
 from pathlib import Path
 from typing import Any
 
-from dotenv import load_dotenv
 import replicate
+from dotenv import load_dotenv
 
 from api.flux import FluxAPI
 
@@ -36,7 +36,7 @@ class PrunaDevAPI(FluxAPI):
             },
         )
         end_time = time.time()
-        
+
         if result:
             self._save_image_from_result(result, save_path)
         else:
@@ -46,4 +46,4 @@ class PrunaDevAPI(FluxAPI):
     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.read())
+            f.write(result.read())

api/replicate.py

@@ -3,8 +3,8 @@ import time
 from pathlib import Path
 from typing import Any
 
-from dotenv import load_dotenv
 import replicate
+from dotenv import load_dotenv
 
 from api.flux import FluxAPI
 

api/replicate_wan.py

@@ -0,0 +1,48 @@
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import replicate
+from dotenv import load_dotenv
+
+from api.flux import FluxAPI
+
+
+class ReplicateAPI(FluxAPI):
+    def __init__(self):
+        load_dotenv()
+        self._api_key = os.getenv("REPLICATE_API_TOKEN")
+        if not self._api_key:
+            raise ValueError("REPLICATE_API_TOKEN not found in environment variables")
+
+    @property
+    def name(self) -> str:
+        return "replicate_go_fast"
+
+    def generate_image(self, prompt: str, save_path: Path) -> float:
+        start_time = time.time()
+        result = replicate.run(
+            "black-forest-labs/flux-dev",
+            input={
+                "seed": 0,
+                "prompt": prompt,
+                "go_fast": True,
+                "guidance": 3.5,
+                "num_outputs": 1,
+                "aspect_ratio": "1:1",
+                "output_format": "png",
+                "num_inference_steps": 28,
+            },
+        )
+        end_time = time.time()
+        if result and len(result) > 0:
+            self._save_image_from_result(result[0], save_path)
+        else:
+            raise Exception("No result returned from Replicate 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)
+        with open(save_path, "wb") as f:
+            f.write(result.read())

api/together.py

@@ -33,10 +33,10 @@ class TogetherAPI(FluxAPI):
             response_format="b64_json",
         )
         end_time = time.time()
-        if result and hasattr(result, 'data') and len(result.data) > 0:
+        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")        
+            raise Exception("No result returned from Together API")
         return end_time - start_time
 
     def _save_image_from_result(self, result: Any, save_path: Path):

benchmark/__init__.py

@@ -7,10 +7,14 @@ from benchmark.hps import HPSPrompts
 from benchmark.parti import PartiPrompts
 
 
-def create_benchmark(benchmark_type: str) -> Type[DrawBenchPrompts | GenAIBenchPrompts | GenEvalPrompts | HPSPrompts | PartiPrompts]:
+def create_benchmark(
+    benchmark_type: str,
+) -> Type[
+    DrawBenchPrompts | GenAIBenchPrompts | GenEvalPrompts | HPSPrompts | PartiPrompts
+]:
     """
     Factory function to create benchmark instances.
-    
+
     Args:
         benchmark_type (str): The type of benchmark to create. Must be one of:
             - "draw_bench"
@@ -18,10 +22,10 @@ def create_benchmark(benchmark_type: str) -> Type[DrawBenchPrompts | GenAIBenchP
             - "geneval"
             - "hps"
             - "parti"
-            
+
     Returns:
         An instance of the requested benchmark implementation
-        
+
     Raises:
         ValueError: If an invalid benchmark type is provided
     """
@@ -32,8 +36,10 @@ def create_benchmark(benchmark_type: str) -> Type[DrawBenchPrompts | GenAIBenchP
         "hps": HPSPrompts,
         "parti": PartiPrompts,
     }
-    
+
     if benchmark_type not in benchmark_map:
-        raise ValueError(f"Invalid benchmark type: {benchmark_type}. Must be one of {list(benchmark_map.keys())}")
-    
+        raise ValueError(
+            f"Invalid benchmark type: {benchmark_type}. Must be one of {list(benchmark_map.keys())}"
+        )
+
     return benchmark_map[benchmark_type]()

benchmark/genai_bench.py

@@ -8,8 +8,8 @@ class GenAIBenchPrompts:
     def __init__(self):
         super().__init__()
         self._download_genai_bench_files()
-        prompts_path = Path('downloads/genai_bench/prompts.txt')
-        with open(prompts_path, 'r') as f:
+        prompts_path = Path("downloads/genai_bench/prompts.txt")
+        with open(prompts_path, "r") as f:
             self.prompts = [line.strip() for line in f if line.strip()]
 
     def __iter__(self) -> Iterator[Tuple[str, Path]]:
@@ -17,13 +17,13 @@ class GenAIBenchPrompts:
             yield prompt, Path(f"{i}.png")
 
     def _download_genai_bench_files(self) -> None:
-        folder_name = Path('downloads/genai_bench')
+        folder_name = Path("downloads/genai_bench")
         folder_name.mkdir(parents=True, exist_ok=True)
         prompts_url = "https://huggingface.co/datasets/zhiqiulin/GenAI-Bench-527/raw/main/prompts.txt"
         prompts_path = folder_name / "prompts.txt"
         if not prompts_path.exists():
             response = requests.get(prompts_url)
-            with open(prompts_path, 'w') as f:
+            with open(prompts_path, "w") as f:
                 f.write(response.text)
 
     @property

benchmark/geneval.py

@@ -9,32 +9,32 @@ class GenEvalPrompts:
     def __init__(self):
         super().__init__()
         self._download_geneval_file()
-        metadata_path = Path('downloads/geneval/evaluation_metadata.jsonl')
+        metadata_path = Path("downloads/geneval/evaluation_metadata.jsonl")
         self.entries: List[Dict[str, Any]] = []
-        with open(metadata_path, 'r') as f:
+        with open(metadata_path, "r") as f:
             for line in f:
                 if line.strip():
                     self.entries.append(json.loads(line))
-    
+
     def __iter__(self) -> Iterator[Tuple[Dict[str, Any], Path]]:
         for i, entry in enumerate(self.entries):
             folder_name = f"{i:05d}"
             yield entry, folder_name
 
     def _download_geneval_file(self) -> None:
-        folder_name = Path('downloads/geneval')
+        folder_name = Path("downloads/geneval")
         folder_name.mkdir(parents=True, exist_ok=True)
         metadata_url = "https://raw.githubusercontent.com/djghosh13/geneval/main/prompts/evaluation_metadata.jsonl"
         metadata_path = folder_name / "evaluation_metadata.jsonl"
         if not metadata_path.exists():
             response = requests.get(metadata_url)
-            with open(metadata_path, 'w') as f:
+            with open(metadata_path, "w") as f:
                 f.write(response.text)
-    
+
     @property
     def name(self) -> str:
         return "geneval"
-    
+
     @property
     def size(self) -> int:
         return len(self.entries)

benchmark/hps.py

@@ -9,13 +9,18 @@ import huggingface_hub
 class HPSPrompts:
     def __init__(self):
         super().__init__()
-        self.hps_prompt_files = ['anime.json', 'concept-art.json', 'paintings.json', 'photo.json']
+        self.hps_prompt_files = [
+            "anime.json",
+            "concept-art.json",
+            "paintings.json",
+            "photo.json",
+        ]
         self._download_benchmark_prompts()
         self.prompts: Dict[str, str] = {}
         self._size = 0
         for file in self.hps_prompt_files:
-            category = file.replace('.json', '')
-            with open(os.path.join('downloads/hps', file), 'r') as f:
+            category = file.replace(".json", "")
+            with open(os.path.join("downloads/hps", file), "r") as f:
                 prompts = json.load(f)
                 for i, prompt in enumerate(prompts):
                     if i == 100:
@@ -23,24 +28,26 @@ class HPSPrompts:
                     filename = f"{category}_{i:03d}.png"
                     self.prompts[filename] = prompt
                     self._size += 1
-    
+
     def __iter__(self) -> Iterator[Tuple[str, Path]]:
         for filename, prompt in self.prompts.items():
             yield prompt, Path(filename)
-    
+
     @property
     def name(self) -> str:
         return "hps"
-    
+
     @property
     def size(self) -> int:
         return self._size
 
     def _download_benchmark_prompts(self) -> None:
-        folder_name = Path('downloads/hps')
+        folder_name = Path("downloads/hps")
         folder_name.mkdir(parents=True, exist_ok=True)
         for file in self.hps_prompt_files:
-            file_name = huggingface_hub.hf_hub_download("zhwang/HPDv2", file, subfolder="benchmark", repo_type="dataset")
+            file_name = huggingface_hub.hf_hub_download(
+                "zhwang/HPDv2", file, subfolder="benchmark", repo_type="dataset"
+            )
             if not os.path.exists(os.path.join(folder_name, file)):
                 os.symlink(file_name, os.path.join(folder_name, file))
 

benchmark/parti.py

@@ -14,11 +14,11 @@ class PartiPrompts:
     def __iter__(self) -> Iterator[Tuple[str, Path]]:
         for i, prompt in enumerate(self.prompts):
             yield prompt, Path(f"{i}.png")
-    
+
     @property
     def name(self) -> str:
         return "parti"
-    
+
     @property
     def size(self) -> int:
         return len(self.prompts)

evaluate.py

@@ -1,59 +1,65 @@
 import argparse
 import json
+import warnings
 from pathlib import Path
 from typing import Dict
-import warnings
 
-from benchmark import create_benchmark
-from benchmark.metrics import create_metric
 import numpy as np
 from PIL import Image
 from tqdm import tqdm
 
+from benchmark import create_benchmark
+from benchmark.metrics import create_metric
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 
 
-def evaluate_benchmark(benchmark_type: str, api_type: str, images_dir: Path = Path("images")) -> Dict:
+def evaluate_benchmark(
+    benchmark_type: str, api_type: str, images_dir: Path = Path("images")
+) -> Dict:
     """
     Evaluate a benchmark's images using its specific metrics.
-    
+
     Args:
         benchmark_type (str): Type of benchmark to evaluate
         api_type (str): Type of API used to generate images
         images_dir (Path): Base directory containing generated images
-        
+
     Returns:
         Dict containing evaluation results
     """
     benchmark = create_benchmark(benchmark_type)
-    
+
     benchmark_dir = images_dir / api_type / benchmark_type
     metadata_file = benchmark_dir / "metadata.jsonl"
-    
+
     if not metadata_file.exists():
-        raise FileNotFoundError(f"No metadata file found for {api_type}/{benchmark_type}. Please run sample.py first.")
-    
+        raise FileNotFoundError(
+            f"No metadata file found for {api_type}/{benchmark_type}. Please run sample.py first."
+        )
+
     metadata = []
     with open(metadata_file, "r") as f:
         for line in f:
             metadata.append(json.loads(line))
-    
-    metrics = {metric_type: create_metric(metric_type) for metric_type in benchmark.metrics}
-    
+
+    metrics = {
+        metric_type: create_metric(metric_type) for metric_type in benchmark.metrics
+    }
+
     results = {
         "api": api_type,
         "benchmark": benchmark_type,
         "metrics": {metric: 0.0 for metric in benchmark.metrics},
-        "total_images": len(metadata)
+        "total_images": len(metadata),
     }
     inference_times = []
-    
+
     for entry in tqdm(metadata):
         image_path = benchmark_dir / entry["filepath"]
         if not image_path.exists():
             continue
-        
+
         for metric_type, metric in metrics.items():
             try:
                 if metric_type == "vqa":
@@ -64,26 +70,30 @@ def evaluate_benchmark(benchmark_type: str, api_type: str, images_dir: Path = Pa
                 results["metrics"][metric_type] += score[metric_type]
             except Exception as e:
                 print(f"Error computing {metric_type} for {image_path}: {str(e)}")
-        
+
         inference_times.append(entry["inference_time"])
-    
+
     for metric in results["metrics"]:
         results["metrics"][metric] /= len(metadata)
     results["median_inference_time"] = np.median(inference_times).item()
-    
+
     return results
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Evaluate generated images using benchmark-specific metrics")
+    parser = argparse.ArgumentParser(
+        description="Evaluate generated images using benchmark-specific metrics"
+    )
     parser.add_argument("api_type", help="Type of API to evaluate")
-    parser.add_argument("benchmarks", nargs="+", help="List of benchmark types to evaluate")
-    
+    parser.add_argument(
+        "benchmarks", nargs="+", help="List of benchmark types to evaluate"
+    )
+
     args = parser.parse_args()
-    
+
     results_dir = Path("evaluation_results")
     results_dir.mkdir(exist_ok=True)
-    
+
     results_file = results_dir / f"{args.api_type}.jsonl"
     existing_results = set()
 
@@ -97,15 +107,15 @@ def main():
         if benchmark_type in existing_results:
             print(f"Skipping {args.api_type}/{benchmark_type} - already evaluated")
             continue
-            
+
         try:
             print(f"Evaluating {args.api_type}/{benchmark_type}")
             results = evaluate_benchmark(benchmark_type, args.api_type)
-            
+
             # Append results to file
             with open(results_file, "a") as f:
                 f.write(json.dumps(results) + "\n")
-                
+
         except Exception as e:
             print(f"Error evaluating {args.api_type}/{benchmark_type}: {str(e)}")