fix: preset prompt persistence

app_local.py

@@ -7,7 +7,6 @@ from PIL import Image
 from diffusers import QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler
 from diffusers.utils import is_xformers_available
 from presets import PRESETS, get_preset_choices, get_preset_info, update_preset_prompt
-
 import os
 import sys
 import re
@@ -16,6 +15,8 @@ import math
 import json  # Added json import
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import logging
+from copy import deepcopy
+
 #############################
 os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
 os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
@@ -26,14 +27,12 @@ logging.basicConfig(
     handlers=[logging.StreamHandler(sys.stdout)]
 )
 logger = logging.getLogger(__name__)
-
 # Model configuration
 REWRITER_MODEL = "Qwen/Qwen1.5-4B-Chat"  # Upgraded to 4B for better JSON handling
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 LOC = os.getenv("QIE")
-
 # Quantization configuration
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -41,7 +40,6 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4",
     bnb_4bit_use_double_quant=True
 )
-
 rewriter_model = AutoModelForCausalLM.from_pretrained(
     REWRITER_MODEL,
     torch_dtype=dtype,
@@ -49,10 +47,16 @@ rewriter_model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config,
 )
 
+def get_fresh_presets():
+    """Return a fresh copy of presets to avoid persistence across users"""
+    return deepcopy(PRESETS)
+
+# Store original presets for reference
+ORIGINAL_PRESETS = deepcopy(PRESETS)
+
 # Preload enhancement model at startup
 print("🔄 Loading prompt enhancement model...")
 rewriter_tokenizer = AutoTokenizer.from_pretrained(REWRITER_MODEL)
-
 print("✅ Enhancement model loaded and ready!")
 
 SYSTEM_PROMPT_EDIT = '''
@@ -242,69 +246,78 @@ try:
 except Exception as e:
     print(f"VAE Slicing Failed: {e}")
     
-  
+
 def toggle_output_count(preset_type):
     """Control output count slider interactivity and show/hide preset editor"""
-    if preset_type and preset_type in PRESETS:
+    if preset_type and preset_type in ORIGINAL_PRESETS:
         # When preset is selected, disable manual output count and show editor
-        preset = PRESETS[preset_type]
+        preset = ORIGINAL_PRESETS[preset_type]
         prompts = preset["prompts"][:4]  # Get up to 4 prompts
         # Pad prompts to 4 items if needed
         while len(prompts) < 4:
             prompts.append("")
         return (
-            gr.Group(visible=True), 
+            gr.Group(visible=True),
             gr.Slider(interactive=False, value=len([p for p in prompts if p.strip()])),  # Count non-empty prompts
             prompts[0], prompts[1], prompts[2], prompts[3]  # Populate preset prompts
         )
     else:
         # When no preset is selected, enable manual output count and hide editor
         return (
-            gr.Group(visible=False), 
+            gr.Group(visible=False),
             gr.Slider(interactive=True),  # Enable slider
             "", "", "", ""  # Clear preset prompts
         )
-### FIXED ### - Updated prompt preview to show only non-empty prompts
+
 def update_prompt_preview(preset_type, base_prompt):
     """Update the prompt preview display based on selected preset and base prompt"""
-    if preset_type and preset_type in PRESETS:
-        preset = PRESETS[preset_type]
+    if preset_type and preset_type in ORIGINAL_PRESETS:
+        preset = ORIGINAL_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."
 
 def update_preset_prompt_textbox(preset_type, prompt_1, prompt_2, prompt_3, prompt_4):
-    """Update preset prompts based on user input"""
-    if preset_type and preset_type in PRESETS:
-        # Update each prompt in the preset
+    """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(PRESETS[preset_type]["prompts"]):
-                PRESETS[preset_type]["prompts"][i] = new_prompt.strip()
+            if i < len(working_presets[preset_type]["prompts"]):
+                working_presets[preset_type]["prompts"][i] = new_prompt.strip()
             else:
-                PRESETS[preset_type]["prompts"].append(new_prompt.strip())
-        
-        # Count non-empty prompts for the slider value
-        non_empty_count = len([p for p in new_prompts if p.strip()])
-        non_empty_count = max(1, min(4, non_empty_count))  # Ensure between 1-4
-        
+                working_presets[preset_type]["prompts"].append(new_prompt.strip())
         # Return updated preset info for preview
-        return update_prompt_preview(preset_type, "your subject")  # Use placeholder 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_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."
 
 @spaces.GPU()
 def infer(
@@ -316,10 +329,10 @@ def infer(
     num_inference_steps=4,
     rewrite_prompt=True,
     num_images_per_prompt=1,
-    preset_type=None,  # New parameter for presets
+    preset_type=None,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Image editing endpoint with optimized prompt handling"""
+    """Image editing endpoint with optimized prompt handling - now uses fresh presets"""
     # Resize image to max 1024px on longest side
     def resize_image(pil_image, max_size=1024):
         """Resize image to maximum dimension of 1024px while maintaining aspect ratio"""
@@ -341,7 +354,7 @@ def infer(
         except Exception as e:
             print(f"⚠️ Image resize failed: {e}")
             return pil_image  # Return original if resize fails
-    
+
     # Add noise function for batch variation
     def add_noise_to_image(pil_image, noise_level=0.001):
         """Add slight noise to image to create variation in outputs"""
@@ -360,17 +373,19 @@ def infer(
             print(f"Warning: Could not add noise to image: {e}")
             return pil_image  # Return original if noise addition fails
 
+    # Get fresh presets for this session
+    session_presets = get_fresh_presets()
+    
     # Resize input image first
     image = resize_image(image, max_size=1024)
     original_prompt = prompt
     prompt_info = ""
     
     # Handle preset batch generation
-    if preset_type and preset_type in PRESETS:
-        preset = PRESETS[preset_type]
+    if preset_type and preset_type in session_presets:
+        preset = session_presets[preset_type]
         # Filter out empty prompts
         non_empty_preset_prompts = [p for p in preset["prompts"] if p.strip()]
-        
         if non_empty_preset_prompts:
             batch_prompts = [f"{original_prompt}, {preset_prompt}" for preset_prompt in non_empty_preset_prompts]
             num_images_per_prompt = len(non_empty_preset_prompts)  # Use actual count of non-empty prompts
@@ -395,7 +410,6 @@ def infer(
             )
     else:
         batch_prompts = [prompt]  # Single prompt in list
-        
         # Handle regular prompt rewriting
         if rewrite_prompt:
             try:
@@ -435,25 +449,20 @@ def infer(
     
     # Set base seed for reproducibility
     base_seed = seed if not randomize_seed else random.randint(0, MAX_SEED)
-    
     try:
         edited_images = []
-        
         # Generate images for each prompt in the batch
         for i, current_prompt in enumerate(batch_prompts):
             # Create unique seed for each image
             generator = torch.Generator(device=device).manual_seed(base_seed + i*1000)
-            
             # Add slight noise to the image for variation (except for first image to maintain base)
-            if i == 0 and len(batch_prompts) == 1:
+            if i == 0 and len(batch_prompts) > 1:
                 input_image = image
             else:
                 input_image = add_noise_to_image(image, noise_level=0.001 + i*0.003)
-            
             # Slightly vary guidance scale for each image
             varied_guidance = true_guidance_scale + random.uniform(-0.1, 0.1)
             varied_guidance = max(1.0, min(10.0, varied_guidance))
-            
             # Generate single image
             result = pipe(
                 image=input_image,
@@ -465,14 +474,11 @@ def infer(
                 num_images_per_prompt=1
             ).images
             edited_images.extend(result)
-            
             print(f"Generated image {i+1}/{len(batch_prompts)} with prompt: {current_prompt[:75]}...")
-        
         # Clear cache after generation
         # if device == "cuda":
         #     torch.cuda.empty_cache()
         #     gc.collect()
-            
         return edited_images, base_seed, prompt_info
     except Exception as e:
         # Clear cache on error
@@ -488,7 +494,8 @@ def infer(
         )
 
 with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
-    preset_prompts_state = gr.State(value=["", "", "", ""])
+    preset_prompts_state = gr.State(value=[])
+    # preset_prompts_state = gr.State(value=["", "", "", ""])
 
     gr.Markdown("""
     <div style="text-align: center; background: linear-gradient(to right, #3a7bd5, #00d2ff); color: white; padding: 20px; border-radius: 8px;">
@@ -602,39 +609,38 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
                 "Prompt details will appear after generation. Ability to edit Preset Prompts on the fly will be implemented shortly.</div>"
             )
 
-    # Function to show preset editor when a preset is selected
+    # Fix the show_preset_editor function to use ORIGINAL_PRESETS:
     def show_preset_editor(preset_type):
-        if preset_type and preset_type in PRESETS:
-            preset = PRESETS[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]
         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 PRESETS:
+        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)
 
-    # Set up preset dropdown change handler - now controls output count slider
+    # 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]
     )
     
-    # Set up prompt preview updates
     preset_dropdown.change(
         fn=update_prompt_preview,
         inputs=[preset_dropdown, prompt],
         outputs=prompt_preview
     )
-    
-    # Set up handlers for preset editor changes
+
     preset_prompt_1.change(
         fn=update_preset_count,
         inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
@@ -655,14 +661,13 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
         inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],
         outputs=num_images_per_prompt
     )
-
+    
     prompt.change(
         fn=update_prompt_preview,
         inputs=[preset_dropdown, prompt],
         outputs=prompt_preview
     )
     
-    # Set up preset editor update button
     update_preset_button.click(
         fn=update_preset_prompt_textbox,
         inputs=[preset_dropdown, preset_prompt_1, preset_prompt_2, preset_prompt_3, preset_prompt_4],