Update app.py

app.py

@@ -12,6 +12,8 @@ from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_d
 import copy
 import random
 import time
+import base64
+import tempfile
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
@@ -113,7 +115,18 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
         joint_attention_kwargs={"scale": lora_scale},
         output_type="pil",
     ).images[0]
-    return final_image 
+
+    # Save the image as a downloadable PNG file
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    final_image.save(temp_file.name, "PNG")
+
+    # Convert the image to a base64 string
+    buffered = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    final_image.save(buffered.name, format="PNG")
+    with open(buffered.name, "rb") as f:
+        img_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+    return final_image, temp_file.name, f"data:image/png;base64,{img_base64}"
 
 @spaces.GPU(duration=30)
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
@@ -154,9 +167,8 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             seed = random.randint(0, MAX_SEED)
             
     if(image_input is not None):
-        
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
-        yield final_image, seed, gr.update(visible=False)
+        final_image, file_path, base64_str = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
+        yield final_image, seed, file_path, base64_str, gr.update(visible=False)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
@@ -167,9 +179,15 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             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 image, seed, None, None, gr.update(value=progress_bar, visible=True)
             
-        yield final_image, seed, gr.update(value=progress_bar, visible=False)
+        # Save the final image and encode to Base64
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        final_image.save(temp_file.name, "PNG")
+        with open(temp_file.name, "rb") as f:
+            img_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+        yield final_image, seed, temp_file.name, f"data:image/png;base64,{img_base64}", gr.update(value=progress_bar, visible=False)
         
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
@@ -218,7 +236,7 @@ def add_custom_lora(custom_lora):
                 <img src="{image}" />
                 <div>
                     <h3>{title}</h3>
-                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
+                    <small>{"Using: <code><b"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
                 </div>
               </div>
             </div>
@@ -295,6 +313,8 @@ with gr.Blocks(theme=gr.themes.Soft(font=font), css=css, delete_cache=(60, 60))
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress",visible=False)
             result = gr.Image(label="Generated Image")
+            download_link = gr.File(label="Download Image")
+            base64_output = gr.Textbox(label="Base64 Encoded Image")
 
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
@@ -312,29 +332,4 @@ with gr.Blocks(theme=gr.themes.Soft(font=font), css=css, delete_cache=(60, 60))
                 
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
-
-    gallery.select(
-        update_selection,
-        inputs=[width, height],
-        outputs=[prompt, selected_info, selected_index, width, height]
-    )
-    custom_lora.input(
-        add_custom_lora,
-        inputs=[custom_lora],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
-    )
-    custom_lora_button.click(
-        remove_custom_lora,
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
-    )
-    gr.on(
-        triggers=[generate_button.click, prompt.submit],
-        fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
-        outputs=[result, seed, progress_bar]
-    )
-
-app.queue()
-app.launch()
+                    seed = gr.Slider(label="Seed", minimum=0, maximum