refactor: streamline image generation process and improve metadata handling in generate_images function

sample.py

@@ -33,32 +33,39 @@ def generate_images(api_type: str, benchmarks: List[str]):
                         entry = json.loads(line)
                         existing_metadata[entry["filepath"]] = entry
 
-            for metadata, folder_name in tqdm(
-                benchmark, desc=f"Generating images for {benchmark_type}", leave=False
-            ):
-                sample_path = benchmark_dir / folder_name
-                samples_path = sample_path / "samples"
-                samples_path.mkdir(parents=True, exist_ok=True)
-                image_path = samples_path / "0000.png"
-
-                if image_path.exists():
-                    continue
-
-                try:
-                    inference_time = api.generate_image(metadata["prompt"], image_path)
-
-                    metadata_entry = {
-                        "filepath": str(image_path),
-                        "prompt": metadata["prompt"],
-                        "inference_time": inference_time,
-                    }
-
-                    existing_metadata[str(image_path)] = metadata_entry
-
-                except Exception as e:
-                    print(f"\nError generating image for prompt: {metadata['prompt']}")
-                    print(f"Error: {str(e)}")
-                    continue
+            with open(metadata_file, "w") as f:
+                for metadata, folder_name in tqdm(
+                    benchmark,
+                    desc=f"Generating images for {benchmark_type}",
+                    leave=False,
+                ):
+                    sample_path = benchmark_dir / folder_name
+                    samples_path = sample_path / "samples"
+                    samples_path.mkdir(parents=True, exist_ok=True)
+                    image_path = samples_path / "0000.png"
+
+                    if image_path.exists():
+                        continue
+
+                    try:
+                        inference_time = api.generate_image(
+                            metadata["prompt"], image_path
+                        )
+
+                        metadata_entry = {
+                            "filepath": str(image_path),
+                            "prompt": metadata["prompt"],
+                            "inference_time": inference_time,
+                        }
+
+                        f.write(json.dumps(metadata_entry) + "\n")
+
+                    except Exception as e:
+                        print(
+                            f"\nError generating image for prompt: {metadata['prompt']}"
+                        )
+                        print(f"Error: {str(e)}")
+                        continue
         else:
             benchmark_dir = api_dir / benchmark_type
             benchmark_dir.mkdir(parents=True, exist_ok=True)
@@ -71,33 +78,35 @@ def generate_images(api_type: str, benchmarks: List[str]):
                         entry = json.loads(line)
                         existing_metadata[entry["filepath"]] = entry
 
-            for prompt, image_path in tqdm(
-                benchmark, desc=f"Generating images for {benchmark_type}", leave=False
-            ):
-                full_image_path = benchmark_dir / image_path
-
-                if full_image_path.exists():
-                    continue
-
-                try:
-                    inference_time = api.generate_image(prompt, full_image_path)
-
-                    metadata_entry = {
-                        "filepath": str(image_path),
-                        "prompt": prompt,
-                        "inference_time": inference_time,
-                    }
-
-                    existing_metadata[str(image_path)] = metadata_entry
-
-                except Exception as e:
-                    print(f"\nError generating image for prompt: {prompt}")
-                    print(f"Error: {str(e)}")
-                    continue
-
             with open(metadata_file, "w") as f:
-                for entry in existing_metadata.values():
-                    f.write(json.dumps(entry) + "\n")
+                for prompt, image_path in tqdm(
+                    benchmark,
+                    desc=f"Generating images for {benchmark_type}",
+                    leave=False,
+                ):
+                    if image_path in existing_metadata:
+                        continue
+
+                    full_image_path = benchmark_dir / image_path
+
+                    if full_image_path.exists():
+                        continue
+
+                    try:
+                        inference_time = api.generate_image(prompt, full_image_path)
+
+                        metadata_entry = {
+                            "filepath": str(image_path),
+                            "prompt": prompt,
+                            "inference_time": inference_time,
+                        }
+
+                        f.write(json.dumps(metadata_entry) + "\n")
+
+                    except Exception as e:
+                        print(f"\nError generating image for prompt: {prompt}")
+                        print(f"Error: {str(e)}")
+                        continue
 
 
 def main():