Update app.py

app.py

@@ -166,7 +166,7 @@ class APIError(Exception):
     pass
 
 
-def create_request(prompt, image, guidance_scale=5.0, image_guidance_scale=4.0, seed=-1):
+def create_request(prompt, image, guidance_scale=5.0, image_guidance_scale=4.0, seed=-1, refine_strength=0.3):
     """
     Create an image editing request to the API.
     
@@ -176,6 +176,7 @@ def create_request(prompt, image, guidance_scale=5.0, image_guidance_scale=4.0,
         guidance_scale (float): Strength of instruction following
         image_guidance_scale (float): Strength of image preservation
         seed (int): Seed for reproducibility, -1 for random
+        refine_strength (float): Strength of refinement
         
     Returns:
         tuple: (task_id, seed) - Task ID if successful and the seed used
@@ -230,7 +231,8 @@ def create_request(prompt, image, guidance_scale=5.0, image_guidance_scale=4.0,
             "custom_params": {
                 "sample_steps": 28,
                 "guidance_scale": guidance_scale,
-                "image_guidance_scale": image_guidance_scale
+                "image_guidance_scale": image_guidance_scale,
+                "refine_strength": refine_strength
             },
         },
         "version": API_VERSION,
@@ -620,13 +622,25 @@ def create_ui():
                             label="Image Preservation Strength"
                         )
                     
-                    seed = gr.Number(
-                        label="Seed (use -1 for random)",
-                        value=82706,
-                        precision=0
-                    )
+                    with gr.Row():
+                        refine_strength = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.01,
+                            value=0.3,
+                            label="Refinement Strength"
+                        )
+                        seed = gr.Number(
+                            label="Seed (use -1 for random)",
+                            value=82706,
+                            precision=0
+                        )
                     
-
+                    gr.Markdown("""
+                    <div style="padding: 8px; margin: 15px 0; background-color: #FFF3CD; border-left: 5px solid #FFDD57; color: #856404;">
+                        <strong>Refinement Strength:</strong> Controls balance between editing and refinement. Higher values enhance details but may reduce editing strength. Set to 0 to disable refinement.
+                    </div>
+                    """)
                 
                 progress = gr.Progress(track_tqdm=False)
 
@@ -663,7 +677,7 @@ def create_ui():
                 return instruction
 
         # Generate function with progress updates
-        def generate_with_progress(image, instruction, seed, guidance_scale, image_guidance_scale, progress=gr.Progress()):
+        def generate_with_progress(image, instruction, seed, guidance_scale, image_guidance_scale, refine_strength, progress=gr.Progress()):
             logger.info(f"Starting image generation with instruction='{instruction[:50]}...', seed={seed}")
             
             try:
@@ -682,7 +696,8 @@ def create_ui():
                     image=image,
                     guidance_scale=guidance_scale,
                     image_guidance_scale=image_guidance_scale,
-                    seed=seed
+                    seed=seed,
+                    refine_strength=refine_strength
                 )
 
                 # Poll for results
@@ -735,6 +750,7 @@ def create_ui():
                                 "model": API_MODEL_NAME,
                                 "guidance_scale": str(guidance_scale),
                                 "image_guidance_scale": str(image_guidance_scale),
+                                "refine_strength": str(refine_strength),
                                 "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                                 "generated_by": "HiDream-E1-Full Editor"
                             }
@@ -748,6 +764,7 @@ def create_ui():
                                 "seed": used_seed,
                                 "guidance_scale": guidance_scale,
                                 "image_guidance_scale": image_guidance_scale,
+                                "refine_strength": refine_strength,
                                 "generated_at": time.strftime("%Y-%m-%d %H:%M:%S")
                             }
 
@@ -781,12 +798,12 @@ def create_ui():
 
         generate_btn.click(
             fn=generate_with_progress,
-            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale],
+            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale, refine_strength],
             outputs=[output_image, image_info]
         )
 
         # Define a combined function to refine instruction and then generate image
-        def refine_and_generate(image, instruction, seed, guidance_scale, image_guidance_scale, progress=gr.Progress()):
+        def refine_and_generate(image, instruction, seed, guidance_scale, image_guidance_scale, refine_strength, progress=gr.Progress()):
             try:
                 # First refine the instruction
                 if not image or not instruction:
@@ -804,7 +821,7 @@ def create_ui():
                 
                 # Then generate with the refined instruction
                 progress(0.2, desc="Instruction refined, generating image...")
-                generated_image, image_info = generate_with_progress(image, refined_instruction, seed, guidance_scale, image_guidance_scale, progress)
+                generated_image, image_info = generate_with_progress(image, refined_instruction, seed, guidance_scale, image_guidance_scale, refine_strength, progress)
                 return generated_image, image_info, refined_instruction
             except Exception as e:
                 logger.error(f"Error in refine_and_generate: {str(e)}")
@@ -815,14 +832,14 @@ def create_ui():
         # Examples
         gr.Examples(
             examples=[
-                ["assets/test_1.png", "convert the image into a Ghibli style",82706, 5, 4],
-                ["assets/test_1.png", "change the image into Disney Pixar style",82706, 5, 4],
-                ["assets/test_1.png", "add a sunglasses to the girl",82706, 5, 2],
-                ["assets/test_2.jpg", "convert this image into a ink sketch image",82706, 5, 2],
-                ["assets/test_2.jpg", "add butterfly",82706, 5, 2],
-                ["assets/test_2.jpg", "remove the wooden sign",82706, 5, 2],
+                ["assets/test_1.png", "convert the image into a Ghibli style",82706, 5, 4, 0.3],
+                ["assets/test_1.png", "change the image into Disney Pixar style",82706, 5, 4, 0.3],
+                ["assets/test_1.png", "add a sunglasses to the girl",82706, 5, 2, 0.3],
+                ["assets/test_2.jpg", "convert this image into a ink sketch image",82706, 5, 2, 0.3],
+                ["assets/test_2.jpg", "add butterfly",82706, 5, 2, 0.3],
+                ["assets/test_2.jpg", "remove the wooden sign",82706, 5, 2, 0.3],
             ],
-            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale],
+            inputs=[input_image, instruction, seed, guidance_scale, image_guidance_scale, refine_strength],
             outputs=[output_image, image_info, instruction],
             fn=refine_and_generate,
             cache_examples=True,