Update app.py

app.py

@@ -152,33 +152,53 @@ def gen(prompt: str, system_prompt: str, token: str, accountId: str):
     images = []
     successful_prompts = []
     failed_prompts = []
-    target_images = 10  # Jumlah gambar yang diinginkan
+    target_images = 10
+    max_attempts_per_image = 3  # Maksimal retry per gambar individual
     
-    while len(images) < target_images:
-        # Generate batch of prompts
-        optimized_prompts = generate_multiple_prompts(prompt, system_prompt)
+    # Generate initial batch of prompts
+    optimized_prompts = generate_multiple_prompts(prompt, system_prompt)
+    
+    # Proses setiap target gambar
+    for i in range(target_images):
+        attempts = 0
+        success = False
         
-        # Try each prompt until we get enough images
-        for opt_prompt in optimized_prompts:
-            if len(images) >= target_images:
-                break
+        # Retry untuk gambar individual ini
+        while attempts < max_attempts_per_image and not success:
+            try:
+                # Ambil prompt dari batch, atau generate baru jika habis
+                if i + attempts >= len(optimized_prompts):
+                    optimized_prompts.extend(generate_multiple_prompts(prompt, system_prompt))
                 
-            img, used_prompt = generate_image_with_fallback(opt_prompt)
-            if img is not None:
-                images.append(img)
-                successful_prompts.append(used_prompt)
-            else:
-                failed_prompts.append(used_prompt)
-                # Jika gagal, generate prompt baru dan coba lagi
-                continue
+                current_prompt = optimized_prompts[i + attempts]
+                img, used_prompt = generate_image_with_fallback(current_prompt)
+                
+                if img is not None:
+                    images.append(img)
+                    successful_prompts.append(used_prompt)
+                    success = True
+                else:
+                    failed_prompts.append(used_prompt)
+                    attempts += 1
+            except Exception as e:
+                failed_prompts.append(f"{current_prompt} (Error: {str(e)})")
+                attempts += 1
+        
+        # Jika masih gagal setelah max attempts
+        if not success:
+            failed_prompts.append(f"Failed after {max_attempts_per_image} attempts for image #{i+1}")
     
     # Prepare output text
-    text = f"Original Prompt: {prompt}\n\nSuccessful Generations ({len(successful_prompts)}):\n"
+    text = f"Original Prompt: {prompt}\n\n"
+    text += f"Successfully Generated: {len(images)} images\n"
+    text += f"Failed Attempts: {len(failed_prompts)}\n\n"
+    
+    text += "Successful Prompts:\n"
     for i, p in enumerate(successful_prompts, 1):
         text += f"\n{i}. {p}"
     
     if failed_prompts:
-        text += f"\n\nFailed Attempts ({len(failed_prompts)}):\n"
+        text += "\n\nFailed Prompts:\n"
         for i, p in enumerate(failed_prompts, 1):
             text += f"\n{i}. {p}"