Update app_local.py

app_local.py

@@ -55,9 +55,9 @@ def get_fresh_presets():
 ORIGINAL_PRESETS = deepcopy(PRESETS)
 
 # Preload enhancement model at startup
-print("🔄 Loading prompt enhancement model...")
+logger.info("🔄 Loading prompt enhancement model...")
 rewriter_tokenizer = AutoTokenizer.from_pretrained(REWRITER_MODEL)
-print("✅ Enhancement model loaded and ready!")
+logger.info("✅ Enhancement model loaded and ready!")
 
 SYSTEM_PROMPT_EDIT = '''
 # Edit Instruction Rewriter
@@ -113,7 +113,7 @@ def extract_json_response(model_output: str) -> str:
         end_idx = model_output.rfind('}')
         # Fix the condition - check if brackets were found
         if start_idx == -1 or end_idx == -1 or start_idx >= end_idx:
-            print(f"No valid JSON structure found in output. Start: {start_idx}, End: {end_idx}")
+            logger.warning(f"No valid JSON structure found in output. Start: {start_idx}, End: {end_idx}")
             return None
         # Expand to the full object including outer braces
         end_idx += 1  # Include the closing brace
@@ -153,8 +153,8 @@ def extract_json_response(model_output: str) -> str:
         if str_values:
             return str_values[0].strip()
     except Exception as e:
-        print(f"JSON parse error: {str(e)}")
-        print(f"Model output was: {model_output}")
+        logger.warning(f"JSON parse error: {str(e)}")
+        logger.warning(f"Model output was: {model_output}")
     return None
 
 def polish_prompt(original_prompt: str) -> str:
@@ -186,8 +186,8 @@ def polish_prompt(original_prompt: str) -> str:
         generated_ids[0][model_inputs.input_ids.shape[1]:],
         skip_special_tokens=True
     ).strip()
-    print(f"Original Prompt: {original_prompt}")
-    print(f"Model raw output: {enhanced}")  # Debug logging
+    logger.info(f"Original Prompt: {original_prompt}")
+    logger.info(f"Model raw output: {enhanced}")  # Debug logging
     # Try to extract JSON content
     rewritten_prompt = extract_json_response(enhanced)
     if rewritten_prompt:
@@ -256,7 +256,7 @@ pipe.fuse_lora()
 try:
     pipe.enable_vae_slicing()
 except Exception as e:
-    print(f"VAE Slicing Failed: {e}")
+    logger.info(f"VAE Slicing Failed: {e}")
     
 
 def toggle_output_count(preset_type):
@@ -297,39 +297,26 @@ def update_prompt_preview(preset_type, base_prompt):
         return preview_text
     else:
         return "Select a preset above to see how your base prompt will be modified for batch generation."
-
-def update_preset_prompt_textbox(preset_type, prompt_1, prompt_2, prompt_3, prompt_4):
-    """Update preset prompts based on user input - now works with session copy"""
-    if preset_type and preset_type in ORIGINAL_PRESETS:
-        # Update each prompt in the preset copy (this won't persist globally)
-        new_prompts = [prompt_1, prompt_2, prompt_3, prompt_4]
-        # Create a working copy for preview purposes
-        working_presets = get_fresh_presets()
-        for i, new_prompt in enumerate(new_prompts):
-            if i < len(working_presets[preset_type]["prompts"]):
-                working_presets[preset_type]["prompts"][i] = new_prompt.strip()
-            else:
-                working_presets[preset_type]["prompts"].append(new_prompt.strip())
-        # Return updated preset info for preview
-        return update_prompt_preview_with_presets(preset_type, "your subject", working_presets)
-    return "Select a preset first to edit its prompts."
+        
+def update_preset_prompt_textbox(preset_type, p1, p2, p3, p4):
+    if preset_type and preset_type in preset_state.value:
+        preset_state.value[preset_type]["prompts"] = [p1, p2, p3, p4]
+        # Re‑render preview with updated data
+        return update_prompt_preview_with_presets(preset_type, prompt.value, preset_state.value)
+    return "Select a preset first."
 
 def update_prompt_preview_with_presets(preset_type, base_prompt, custom_presets):
-    """Update the prompt preview display with custom presets"""
     if preset_type and preset_type in custom_presets:
         preset = custom_presets[preset_type]
         non_empty_prompts = [p for p in preset["prompts"] if p.strip()]
         if not non_empty_prompts:
             return "No prompts defined. Please enter at least one prompt in the editor."
-        preview_text = f"**Preset: {preset_type}**\n\n"
-        preview_text += f"*{preset['description']}*\n\n"
-        preview_text += f"**Generating {len(non_empty_prompts)} image{'s' if len(non_empty_prompts) > 1 else ''}:**\n"
-        for i, preset_prompt in enumerate(non_empty_prompts, 1):
-            full_prompt = f"{base_prompt}, {preset_prompt}"
-            preview_text += f"{i}. {full_prompt}\n"
-        return preview_text
-    else:
-        return "Select a preset above to see how your base prompt will be modified for batch generation."
+        preview = f"**Preset: {preset_type}**\n\n{preset['description']}\n\n"
+        preview += f"**Generating {len(non_empty_prompts)} image{'s' if len(non_empty_prompts)>1 else ''}:**\n"
+        for i, pp in enumerate(non_empty_prompts, 1):
+            preview += f"{i}. {base_prompt}, {pp}\n"
+        return preview
+    return "Select a preset to see the preview."
 
 @spaces.GPU()
 def infer(
@@ -346,6 +333,8 @@ def infer(
 ):
     """Image editing endpoint with optimized prompt handling - now uses fresh presets"""
     # Resize image to max 1024px on longest side
+    session_presets = preset_state.value
+
     def resize_image(pil_image, max_size=1024):
         """Resize image to maximum dimension of 1024px while maintaining aspect ratio"""
         try:
@@ -361,10 +350,10 @@ def infer(
             new_height = int(height * scale)
             # Resize image
             resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
-            print(f"📝 Image resized from {width}x{height} to {new_width}x{new_height}")
+            logger.info(f"📝 Image resized from {width}x{height} to {new_width}x{new_height}")
             return resized_image
         except Exception as e:
-            print(f"⚠️ Image resize failed: {e}")
+            logger.warning(f"⚠️ Image resize failed: {e}")
             return pil_image  # Return original if resize fails
 
     # Add noise function for batch variation
@@ -382,7 +371,7 @@ def infer(
             noisy_array = (noisy_array * 255).astype(np.uint8)
             return Image.fromarray(noisy_array)
         except Exception as e:
-            print(f"Warning: Could not add noise to image: {e}")
+            logger.warning(f"Warning: Could not add noise to image: {e}")
             return pil_image  # Return original if noise addition fails
 
     # Get fresh presets for this session
@@ -409,7 +398,7 @@ def infer(
                 f"<p>Generating {len(non_empty_preset_prompts)} image{'s' if len(non_empty_preset_prompts) > 1 else ''}</p>"
                 f"</div>"
             )
-            print(f"Using preset: {preset_type} with {len(batch_prompts)} variations")
+            logger.info(f"Using preset: {preset_type} with {len(batch_prompts)} variations")
         else:
             # Fallback to manual if no valid prompts
             batch_prompts = [prompt]
@@ -487,7 +476,7 @@ def infer(
                 num_images_per_prompt=2
             ).images
             edited_images.extend(result)
-            print(f"Generated image {i+1}/{len(batch_prompts)} with prompt: {current_prompt}...")
+            logger.info(f"Generated image {i+1}/{len(batch_prompts)} with prompt: {current_prompt}...")
         # Clear cache after generation
         # if device == "cuda":
         #     torch.cuda.empty_cache()
@@ -509,7 +498,8 @@ def infer(
 with gr.Blocks(title="'Qwen Image Edit' Model Playground & Showcase [4-Step Lightning Mode]") as demo:
     preset_prompts_state = gr.State(value=[])
     # preset_prompts_state = gr.State(value=["", "", "", ""])
-
+    preset_state = gr.State(value=get_fresh_presets())
+    
     gr.Markdown("""
     <div style="text-align: center; background: linear-gradient(to right, #3a7bd5, #00d2ff); color: white; padding: 20px; border-radius: 8px;">
         <h1 style="margin-bottom: 5px;">⚡️ Qwen-Image-Edit Lightning</h1>
@@ -626,31 +616,26 @@ with gr.Blocks(title="'Qwen Image Edit' Model Playground & Showcase [4-Step Ligh
                 "Hint: depending on the original image, prompt quality, and complexity, you can often get away with 3 steps, even 2 steps without much loss in quality. </div>"
             )
 
-    # Fix the show_preset_editor function to use ORIGINAL_PRESETS:
     def show_preset_editor(preset_type):
-        if preset_type and preset_type in ORIGINAL_PRESETS:  # Changed from PRESETS to ORIGINAL_PRESETS
-            preset = ORIGINAL_PRESETS[preset_type]
-            prompts = preset["prompts"]
-            # Pad prompts to 4 items if needed
-            while len(prompts) < 4:
-                prompts.append("")
-            return gr.Group(visible=True), prompts[0], prompts[1], prompts[2], prompts[3]
+        if preset_type and preset_type in preset_state.value:
+            preset = preset_state.value[preset_type]
+            prompts = preset["prompts"] + [""] * (4 - len(preset["prompts"]))
+            return gr.Group(visible=True), *prompts[:4]
         return gr.Group(visible=False), "", "", "", ""
-    
-    # Fix the update_preset_count function to use ORIGINAL_PRESETS:
-    def update_preset_count(preset_type, prompt_1, prompt_2, prompt_3, prompt_4):
-        """Update the output count slider based on non-empty preset prompts"""
-        if preset_type and preset_type in ORIGINAL_PRESETS:  # Changed from PRESETS to ORIGINAL_PRESETS
-            non_empty_count = len([p for p in [prompt_1, prompt_2, prompt_3, prompt_4] if p.strip()])
-            return gr.Slider(value=max(1, min(4, non_empty_count)), interactive=False)
-        return gr.Slider(interactive=True, visible=False)
+        
+    def update_preset_count(preset_type, p1, p2, p3, p4):
+        if preset_type and preset_type in preset_state.value:
+            count = len([p for p in (p1,p2,p3,p4) if p.strip()])
+            return gr.Slider(value=max(1, min(4, count)), interactive=False)
+        return gr.Slider(interactive=True)
 
     # Update the preset_dropdown.change handlers to use ORIGINAL_PRESETS
     preset_dropdown.change(
-        fn=toggle_output_count,
-        inputs=preset_dropdown,
-        outputs=[preset_editor, num_images_per_prompt, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4]
+        fn=show_preset_editor,
+        inputs=[preset_dropdown],
+        outputs=[preset_editor, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4]
     )
+
     
     preset_dropdown.change(
         fn=update_prompt_preview,
@@ -658,6 +643,28 @@ with gr.Blocks(title="'Qwen Image Edit' Model Playground & Showcase [4-Step Ligh
         outputs=prompt_preview
     )
 
+    preset_prompt_1.change(
+        fn=update_preset_prompt_textbox,
+        inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
+        outputs=prompt_preview
+    )
+
+    preset_prompt_2.change(
+        fn=update_preset_prompt_textbox,
+        inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
+        outputs=prompt_preview
+    )
+    preset_prompt_3.change(
+        fn=update_preset_prompt_textbox,
+        inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
+        outputs=prompt_preview
+    )
+    preset_prompt_4.change(
+        fn=update_preset_prompt_textbox,
+        inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
+        outputs=prompt_preview
+    )
+
     preset_prompt_1.change(
         fn=update_preset_count,
         inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],