Update app.py

app.py

@@ -31,6 +31,12 @@ longformer_model = LongformerModel.from_pretrained("allenai/longformer-base-4096
 df = pd.read_csv('prompts.csv', header=None)
 prompt_values = df.values.flatten()
 
+# Define the folder path for saving images
+SAVE_PATH = "d:/autosave"
+
+# Ensure the save path exists
+os.makedirs(SAVE_PATH, exist_ok=True)
+
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
@@ -373,35 +379,26 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
                 low_cpu_mem_usage=True,
                 adapter_name=lora_name
             )
-#            if image_input is not None: pipe_i2i = pipe_to_use
-#            else: pipe = pipe_to_use
         print("Loaded LoRAs:", lora_names)
         print("Adapter weights:", lora_weights)
-#        if image_input is not None:
-#            pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
-#        else:
         pipe.set_adapters(lora_names, adapter_weights=lora_weights)
-#    print(pipe.get_active_adapters())
-    # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
-# Generate image
-#    if image_input is not None:
-#        final_image = generate_image_to_image(prompt_mash, steps, cfg_scale, width, height, seed)
-#        yield final_image, seed, gr.update(visible=False)
-#    else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
-# Consume the generator to get the final image
-#        final_image = None
         step_counter = 0
         for image in image_generator:
             step_counter += 1
             final_image = image
             progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
             yield image, seed, gr.update(value=progress_bar, visible=True)
-#            yield final_image, seed, gr.update(value=progress_bar, visible=False)
+
+         # Save the final image if enabled
+        if enable_save and final_image is not None:
+            save_path = os.path.join(SAVE_PATH, f"{seed}.png")
+            final_image.save(save_path)
+            print(f"Image saved to: {save_path}")
 
 run_lora.zerogpu = True
 
@@ -527,6 +524,9 @@ with gr.Blocks(css=css, delete_cache=(240, 240)) as app:
             
         with gr.Column(scale=1,min_width=50):
             randomize_button = gr.Button("🎲", variant="secondary", scale=1, elem_id="random_btn")
+
+        with gr.Row():
+            enable_save = gr.Checkbox(label="Enable Auto-Save", value=False)
                 
 #    with gr.Row(elem_id="loaded_loras"):            
 #        with gr.Column(scale=8):
@@ -598,7 +598,8 @@ with gr.Blocks(css=css, delete_cache=(240, 240)) as app:
     gallery.select(
         update_selection,
         inputs=[selected_indices, loras_state, width, height],
-        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2])
+        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2]
+    )    
     remove_button_1.click(
         remove_lora_1,
         inputs=[selected_indices, loras_state],
@@ -627,7 +628,7 @@ with gr.Blocks(css=css, delete_cache=(240, 240)) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+        inputs=[prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, enable_save],
         outputs=[result, seed, progress_bar]
     )#.then(
      #   fn=lambda x, history: update_history(x, history),