add prompt count

app.py

@@ -49,13 +49,16 @@ def change_image_style(
         edit_prompt = data.get("edit_prompt")
         # replace \n to <br>
         edit_prompt = edit_prompt.replace("\n", "<br>")
-        return resultImageUrl, edit_prompt
+        return resultImageUrl, edit_prompt, ""
     else:
-        raise Exception(f"Error: {response.status_code} - {response.text}")
+        error_msg = f"Error: {response.status_code} - {response.text}"
+        return "", "", error_msg
 
 def generate_image(
     input_image,
     style_images,
+    prompt_count: int,
+    image_count: int,
     extract_model,
     extract_prompt,
     prompt_prefix
@@ -66,20 +69,47 @@ def generate_image(
 
     if style_images is None:
         raise gr.Error(f"Cannot find any style image! Please refer to step 1️⃣")
+    
+    style_size = len(style_images)
+    if style_size < 1:
+        raise gr.Error(f"Please upload at least one style image! Refer to step 1️⃣")
+    
+    total_count = style_size * prompt_count * image_count
+    completed_count = 0
+    success_count = 0
+    error_count = 0
 
     inputImageUrl = prefix + input_image
 
-    result_images = []
+    result_image_list = []
+    error_list = []
     markdownStr = ""
-    for style_image in style_images:
-        if not style_image:
-            raise gr.Error(f"Cannot find any style image! Please refer to step 1️⃣")
-        styleImageUrl = prefix + style_image[0]
-        result_image, edit_prompt = change_image_style(inputImageUrl, styleImageUrl, extract_model, extract_prompt, prompt_prefix)
-        result_images.append(result_image)
-        markdownStr += f"{edit_prompt} <img src='{result_image}' style='zoom: 33%;' />\n"
 
-    return result_images, markdownStr
+    for style_index in range(style_size):
+        style_image = style_images[style_index]
+        if not style_image:
+            error_list.append(f"Style image {style_index + 1} is empty.")
+            continue
+
+        for prompt_index in range(prompt_count):
+            for image_index in range(image_count):
+                completed_count += 1
+                styleImageUrl = prefix + style_image[0]
+                result_image, edit_prompt, error_msg = change_image_style(
+                    inputImageUrl, styleImageUrl, extract_model, extract_prompt, prompt_prefix
+                )
+
+                if result_image:
+                    success_count += 1
+                    result_image_list.append(result_image)
+                    markdownStr += f"{edit_prompt} <img src='{result_image}' style='zoom: 33%;' />\n"
+                else:
+                    error_count += 1
+                    error_list.append(f"Error processing style image {style_index + 1}, prompt {prompt_index + 1}, image {image_index + 1}: {error_msg}")
+
+                status = "processing" if completed_count < total_count else "completed"
+                status_message = f"Status: {status} ({completed_count}/{total_count}) success: {success_count}, error: {error_count}"
+                yield status_message, error_list, result_image_list, markdownStr
 
 
 def swap_to_gallery(images):
@@ -101,6 +131,11 @@ with gr.Blocks() as demo:
             with gr.Column(visible=False) as clear_button:
                 remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
             input_image = gr.Image(label="Input Image", type="filepath", interactive=True)
+            with gr.Row():
+                with gr.Column():
+                    prompt_count = gr.Number(label="Prompt Count", value=4, precision=0)
+                with gr.Column():
+                    image_count = gr.Number(label="Image Count", value=3, precision=0)
             extract_model = gr.Dropdown(choices=extract_model_tag_list, label="Extract Model", value="o3")
             extract_prompt = gr.Textbox(lines=2, label="Extracted Prompt", value="Accurately extract everything from this photo into a prompt especially the hairstyle excluding the backgrounds. I want to replicate it perfectly. need to be a complete paragraph")
             prompt_prefix = gr.Textbox(lines=1, label="Prompt Prefix", value="change the backgrounds to match this:")
@@ -108,6 +143,8 @@ with gr.Blocks() as demo:
             submit = gr.Button("Submit")
 
         with gr.Column():
+            status_text = gr.Textbox(label="Status", show_label=False, interactive=False)
+            error_text = gr.Textbox(label="Error List", show_label=False, interactive=False)
             gallery = gr.Gallery(label="Generated Images")
             markdown = gr.Markdown(label="Generated Images")
 
@@ -116,8 +153,8 @@ with gr.Blocks() as demo:
 
         submit.click(
             fn=generate_image,
-            inputs=[input_image, uploaded_files, extract_model, extract_prompt, prompt_prefix],
-            outputs=[gallery, markdown]
+            inputs=[input_image, uploaded_files, prompt_count, image_count, extract_model, extract_prompt, prompt_prefix],
+            outputs=[status_text, error_text, gallery, markdown]
         )