save details

app.py

@@ -8,6 +8,7 @@ import json
 from PIL import Image
 import torch
 from torchvision import transforms
+import zipfile
 
 from diffusers import FluxFillPipeline, AutoencoderKL
 from PIL import Image
@@ -76,10 +77,9 @@ def calculate_optimal_dimensions(image: Image.Image):
     return width, height
 
 @spaces.GPU(durations=300)
-def infer(edit_images, prompt, lora_model, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     # pipe.enable_xformers_memory_efficient_attention()
 
-
     if lora_model != "None":
         pipe.load_lora_weights(lora_models[lora_model])
         pipe.enable_lora()
@@ -100,6 +100,7 @@ def infer(edit_images, prompt, lora_model, seed=42, randomize_seed=False, guidan
         height=height,
         width=width,
         guidance_scale=guidance_scale,
+        strength=strength,
         num_inference_steps=num_inference_steps,
         generator=torch.Generator(device='cuda').manual_seed(seed),
         # lora_scale=0.75 // not supported in this version
@@ -117,6 +118,35 @@ def download_image(image):
     image.save("output.png", "PNG")
     return "output.png"
 
+def save_details(image, mask, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    if isinstance(mask, np.ndarray):
+        mask = Image.fromarray(mask)
+
+    image.save("saved_image.png", "PNG")
+    mask.save("saved_mask.png", "PNG")
+
+    details = {
+        "prompt": prompt,
+        "lora_model": lora_model,
+        "strength": strength,
+        "seed": seed,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps
+    }
+
+    with open("details.json", "w") as f:
+        json.dump(details, f)
+
+    # Create a ZIP file
+    with zipfile.ZipFile("output.zip", "w") as zipf:
+        zipf.write("saved_image.png")
+        zipf.write("saved_mask.png")
+        zipf.write("details.json")
+
+    return "output.zip"
+
 def set_image_as_inpaint(image):
     return image
 
@@ -203,7 +233,15 @@ with gr.Blocks(css=css) as demo:
                     value=28,
                 )
 
-            # with gr.Row():
+            with gr.Row():
+
+                strength = gr.Slider(
+                    label="Strength",
+                    minimum=0,
+                    maximum=1,
+                    step=0.01,
+                    value=0.85,
+                )
 
             #     width = gr.Slider(
             #         label="width",
@@ -224,12 +262,13 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [edit_image, prompt, lora_model, seed, randomize_seed, guidance_scale, num_inference_steps],
+        inputs = [edit_image, prompt, lora_model, strength, seed, randomize_seed, guidance_scale, num_inference_steps],
         outputs = [result, seed]
     )
 
     download_button = gr.Button("Download Image as PNG")
     set_inpaint_button = gr.Button("Set Image as Inpaint")
+    save_button = gr.Button("Save Details")
 
     download_button.click(
             fn=download_image,
@@ -243,6 +282,12 @@ with gr.Blocks(css=css) as demo:
             outputs=[edit_image]
     )
 
+    save_button.click(
+            fn=save_details,
+            inputs=[result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps],
+            outputs=gr.Text(label="Save Status")
+    )
+
     # edit_image.select(
     #     fn=generate_mask,
     #     inputs=[edit_image, gr.Number(), gr.Number()],