progress

asset-work/generate_multiple.py

@@ -23,25 +23,21 @@ def check_server(server_address: str = "127.0.0.1:8188", timeout: int = 5) -> bo
 
 def count_existing_variations(caption_file: str) -> int:
     """Count existing image variations for a caption file.
-    
+
     Args:
         caption_file: Path to the caption text file
-        
+
     Returns:
-        Number of existing variation directories
+        Number of existing generated images matching the caption name
     """
-    caption_dir = os.path.dirname(os.path.abspath(caption_file))
     caption_name = os.path.splitext(os.path.basename(caption_file))[0]
-    
-    if not os.path.exists(caption_dir):
-        return 0
-    
+    caption_dir = os.path.dirname(caption_file) or "."
+
     count = 0
     for item in os.listdir(caption_dir):
-        item_path = os.path.join(caption_dir, item)
-        if os.path.isdir(item_path) and item.startswith(f"{caption_name}_"):
+        if item.startswith(f"{caption_name}_") and item.endswith("_generated.png"):
             count += 1
-    
+
     return count
 
 
@@ -55,7 +51,7 @@ def generate_multiple(caption_file: str, count: int = 2, server: str = "127.0.0.
         dry_run: If True, validate only without generating
 
     Returns:
-        Tuple of (list of created output directories, number of existing variations)
+        Tuple of (list of created image paths, number of existing variations)
     """
     if not os.path.exists(caption_file):
         print(f"Error: Caption file not found: {caption_file}")
@@ -109,11 +105,10 @@ def generate_multiple(caption_file: str, count: int = 2, server: str = "127.0.0.
         if result.returncode != 0:
             print(f"Warning: Generation {i + 1} of {needed} failed with exit code {result.returncode}")
         else:
-            # Calculate output directory name
-            caption_dir = os.path.dirname(os.path.abspath(caption_file))
+            # Image is saved in current directory by generate_from_caption
             caption_name = os.path.splitext(os.path.basename(caption_file))[0]
-            output_dir = os.path.join(caption_dir, f"{caption_name}_{seed}")
-            output_dirs.append(output_dir)
+            output_path = os.path.join(".", f"{caption_name}_{seed}_generated.png")
+            output_dirs.append(output_path)
 
     return output_dirs, existing
 
@@ -178,16 +173,16 @@ def main():
         sys.exit(1)
     print("✓ Server is running\n")
 
-    output_dirs, existing = generate_multiple(args.caption_file, args.count, args.server)
+    output_paths, existing = generate_multiple(args.caption_file, args.count, args.server)
 
     print(f"\n{'='*60}")
     print("All generations complete!")
     print(f"{'='*60}")
-    if output_dirs:
-        print(f"\nCreated {len(output_dirs)} new output directory(s):")
-        for output_dir in output_dirs:
-            print(f"  - {output_dir}")
-    print(f"\nTotal variations now: {existing + len(output_dirs)}")
+    if output_paths:
+        print(f"\nCreated {len(output_paths)} new image(s):")
+        for output_path in output_paths:
+            print(f"  - {output_path}")
+    print(f"\nTotal variations now: {existing + len(output_paths)}")
 
 
 if __name__ == "__main__":