Update app.py

app.py

@@ -16,7 +16,7 @@ pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef
 MAX_SEED = 2**32-1
 
 @spaces.GPU()
-def generate_image(prompt, width, height, lora_path, trigger_word, steps):
+def generate_image(prompt, width, height, lora_path, trigger_word, steps, custom_hash):
     # Load LoRA weights
     pipe.load_lora_weights(lora_path)
     
@@ -37,10 +37,14 @@ def generate_image(prompt, width, height, lora_path, trigger_word, steps):
         generator=generator,
     ).images[0]
     
-    # Generate a hash for the image
-    image_bytes = image.tobytes()
-    hash_object = hashlib.sha256(image_bytes)
-    image_hash = hash_object.hexdigest()
+    # Generate or use provided hash for the image
+    if not custom_hash:
+        # Generate a hash if custom_hash is not provided
+        image_bytes = image.tobytes()
+        hash_object = hashlib.sha256(image_bytes)
+        image_hash = hash_object.hexdigest()
+    else:
+        image_hash = custom_hash
     
     # Save the image with the hash as filename
     image_path = f"{image_hash}.png"
@@ -48,8 +52,8 @@ def generate_image(prompt, width, height, lora_path, trigger_word, steps):
     
     return image, image_hash
 
-def run_lora(prompt, width, height, lora_path, trigger_word, steps):
-    return generate_image(prompt, width, height, lora_path, trigger_word, steps)
+def run_lora(prompt, width, height, lora_path, trigger_word, steps, custom_hash):
+    return generate_image(prompt, width, height, lora_path, trigger_word, steps, custom_hash)
 
 # Set up the Gradio interface
 with gr.Blocks() as app:
@@ -69,6 +73,9 @@ with gr.Blocks() as app:
     with gr.Row():
         steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=28)
     
+    with gr.Row():
+        custom_hash = gr.Textbox(label="Custom Hash (optional)", placeholder="Leave blank to auto-generate hash")
+    
     generate_button = gr.Button("Generate Image")
     
     output_image = gr.Image(label="Generated Image")
@@ -76,7 +83,7 @@ with gr.Blocks() as app:
     
     generate_button.click(
         fn=run_lora,
-        inputs=[prompt, width, height, lora_path, trigger_word, steps],
+        inputs=[prompt, width, height, lora_path, trigger_word, steps, custom_hash],
         outputs=[output_image, output_hash]
     )