Upload folder using huggingface_hub

app.py

@@ -14,7 +14,7 @@ WORKFLOW_FILENAME = 'sd3.json'
 ROOT_DIR = os.getcwd()  # Root directory of the project
 OUTPUT_DIR = os.path.join(ROOT_DIR, 'output')
 USER_OUTPUTS_DIR = os.path.join(OUTPUT_DIR, 'user_outputs')
-URL = "https://3653-47-216-153-41.ngrok-free.app/prompt"  # Replace with your ngrok URL
+URL = "https://your-ngrok-url.ngrok-free.app/prompt"  # Replace with your ngrok URL
 
 # Create output directories if they don't exist
 os.makedirs(USER_OUTPUTS_DIR, exist_ok=True)
@@ -92,22 +92,28 @@ def monitor_output_images(output_dir, previous_images, timeout=60):
 
 def copy_file_with_retry(src, dst_dir, file_index, retries=5, delay=1):
     dst = os.path.join(dst_dir, f"SD3_{file_index:05d}.png")
+    log_message(f"Attempting to copy file from {src} to {dst}")
     for _ in range(retries):
         try:
             shutil.copy(src, dst)
+            log_message(f"File copied successfully to {dst}")
             return dst
         except PermissionError:
             time.sleep(delay)
+            log_message(f"PermissionError encountered. Retrying...")
+    log_message(f"Failed to copy {src} to {dst} after {retries} retries")
     raise PermissionError(f"Failed to copy {src} to {dst} after {retries} retries")
 
 def zip_files(output_images: List[str], zip_interval: int, zip_folder: str):
     zip_files = []
     for i in range(0, len(output_images), zip_interval):
         zip_filename = os.path.join(zip_folder, f"images_{i//zip_interval + 1}_{time.time_ns()}.zip")
+        log_message(f"Creating zip file {zip_filename}")
         with zipfile.ZipFile(zip_filename, 'w') as zipf:
             for img in output_images[i:i+zip_interval]:
                 zipf.write(img, os.path.basename(img))
         zip_files.append(zip_filename)
+    log_message(f"Zip files created: {zip_files}")
     return zip_files
 
 def process_prompts(prompts_text, negative_prompt_text, user_folder, zip_interval):
@@ -167,6 +173,7 @@ def process_prompts(prompts_text, negative_prompt_text, user_folder, zip_interva
                 continue
 
             new_image_path = os.path.join(OUTPUT_DIR, new_image)
+            log_message(f"New image generated: {new_image_path}")
             try:
                 copied_image_path = copy_file_with_retry(new_image_path, user_folder, file_index)
                 log_message(f"New image generated and copied to user folder: {copied_image_path}")
@@ -251,4 +258,65 @@ def main():
                 with gr.Accordion("Instructions", open=True):
                     gr.Markdown("""
                     **SD3 Batch Imagine Instructions**
-                    -
+                    - Enter your prompts below, one per empty line.
+                    - Enter your negative prompts below, one per line. (Optional)
+                    - Set the zip interval to determine how many images will be included in each zip file.
+                    - Click "Process Prompts" to start generating images.
+                    - Click "Cancel Processing" to stop the current batch run.
+                    - Watch the progress as images are generated in real-time.
+                    - At the end of the process, zip files containing your images will be available for download.
+                    """)
+
+            with gr.Column(scale=2):
+                gr.Markdown("### Enter Prompts")
+                prompts_text = gr.Textbox(lines=20, placeholder="Enter your prompts here, one per empty line.", label="Prompts")
+                negative_prompts_text = gr.Textbox(lines=5, placeholder="Enter your negative prompts here, one per line.", label="Negative Prompts")
+                zip_interval = gr.Number(value=10, label="Zip Interval", precision=0)
+                process_btn = gr.Button("Process Prompts")
+                cancel_btn = gr.Button("Cancel Processing")
+                progress_text = gr.Markdown("Progress")
+                gallery_output = gr.Gallery(label="Generated Images")
+                zip_files_output = gr.Files(label="Zip Files")
+
+            with gr.Column(scale=1):
+                gr.Markdown("### Detailed Logs")
+                logs_output = gr.Textbox(lines=10, interactive=False, label="Logs")
+
+                def generate_user_folder():
+                    user_folder = os.path.normpath(os.path.join(USER_OUTPUTS_DIR, f'SD3{random.randint(1000, 9999)}'))
+                    os.makedirs(user_folder, exist_ok=True)
+                    log_message(f"Generated user folder: {user_folder}")
+                    return user_folder
+
+                def on_click(prompts_text, negative_prompts_text, zip_interval):
+                    reset_cancel_processing_fn()
+                    user_folder = generate_user_folder()
+                    output_images, zip_files_list = [], []
+                    logs = ""
+                    try:
+                        for images, zip_files, log_msg in process_prompts(prompts_text, negative_prompts_text, user_folder, zip_interval):
+                            output_images = images
+                            zip_files_list = zip_files
+                            logs = log_msg
+                            yield images, zip_files_list, logs + log_messages
+                    except Exception as e:
+                        log_message(f"Error during prompt processing: {e}")
+                        logs += f"Error: {e}\n"
+                        yield output_images, zip_files_list, logs + log_messages
+
+                process_btn.click(
+                    fn=on_click,
+                    inputs=[prompts_text, negative_prompts_text, zip_interval],
+                    outputs=[gallery_output, zip_files_output, logs_output]
+                )
+
+                cancel_btn.click(
+                    fn=cancel_processing_fn,
+                    inputs=[],
+                    outputs=[]
+                )
+
+    demo.launch()
+
+if __name__ == "__main__":
+    main()