chore: implement presets in main app (#3)

- chore: implement presets in main app (2bb4b9cb2ae2a8a09632fda67260cea229ad3216)

app_local.py

@@ -6,6 +6,8 @@ import spaces
 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
+
 import os
 import sys
 import re
@@ -85,7 +87,6 @@ Please provide the rewritten instruction in a clean `json` format as:
 }
 '''
 
-
 def extract_json_response(model_output: str) -> str:
     """Extract rewritten instruction from potentially messy JSON output"""
     # Remove code block markers first
@@ -94,19 +95,15 @@ def extract_json_response(model_output: str) -> str:
         # Find the JSON portion in the output
         start_idx = model_output.find('{')
         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}")
             return None
-            
         # Expand to the full object including outer braces
         end_idx += 1  # Include the closing brace
         json_str = model_output[start_idx:end_idx]
-        
         # Handle potential markdown or other formatting
         json_str = json_str.strip()
-        
         # Try to parse JSON directly first
         try:
             data = json.loads(json_str)
@@ -119,7 +116,6 @@ def extract_json_response(model_output: str) -> str:
             json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
             # Try parsing again
             data = json.loads(json_str)
-        
         # Extract rewritten prompt from possible key variations
         possible_keys = [
             "Rewritten", "rewritten", "Rewrited", "rewrited", "Rewrittent",
@@ -128,45 +124,36 @@ def extract_json_response(model_output: str) -> str:
         for key in possible_keys:
             if key in data:
                 return data[key].strip()
-        
         # Try nested path
         if "Response" in data and "Rewritten" in data["Response"]:
             return data["Response"]["Rewritten"].strip()
-        
         # Handle nested JSON objects (additional protection)
         if isinstance(data, dict):
             for value in data.values():
                 if isinstance(value, dict) and "Rewritten" in value:
                     return value["Rewritten"].strip()
-        
         # Try to find any string value that looks like an instruction
         str_values = [v for v in data.values() if isinstance(v, str) and 10 < len(v) < 500]
         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}")
     return None
 
-
 def polish_prompt(original_prompt: str) -> str:
     """Enhanced prompt rewriting using original system prompt with JSON handling"""
-    
     # Format as Qwen chat
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT_EDIT},
         {"role": "user", "content": original_prompt}
     ]
-    
     text = rewriter_tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    
     model_inputs = rewriter_tokenizer(text, return_tensors="pt").to(device)
-    
     with torch.no_grad():
         generated_ids = rewriter_model.generate(
             **model_inputs,
@@ -178,18 +165,14 @@ def polish_prompt(original_prompt: str) -> str:
             no_repeat_ngram_size=3,
             pad_token_id=rewriter_tokenizer.eos_token_id
         )
-    
     # Extract and clean response
     enhanced = rewriter_tokenizer.decode(
         generated_ids[0][model_inputs.input_ids.shape[1]:],
         skip_special_tokens=True
     ).strip()
-    
     print(f"Model raw output: {enhanced}")  # Debug logging
-    
     # Try to extract JSON content
     rewritten_prompt = extract_json_response(enhanced)
-    
     if rewritten_prompt:
         # Clean up remaining artifacts
         rewritten_prompt = re.sub(r'(Replace|Change|Add) "(.*?)"', r'\1 \2', rewritten_prompt)
@@ -205,12 +188,10 @@ def polish_prompt(original_prompt: str) -> str:
                 rewritten_prompt = enhanced
         else:
             rewritten_prompt = enhanced
-        
         # Basic cleanup
         rewritten_prompt = re.sub(r'\s\s+', ' ', rewritten_prompt).strip()
         if ': ' in rewritten_prompt:
             rewritten_prompt = rewritten_prompt.split(': ', 1)[-1].strip()
-        
         return rewritten_prompt[:200] if rewritten_prompt else original_prompt
 
 # Scheduler configuration for Lightning
@@ -231,6 +212,7 @@ scheduler_config = {
     "use_karras_sigmas": False,
 }
 
+
 # Initialize scheduler with Lightning config
 scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
 
@@ -254,15 +236,7 @@ if is_xformers_available():
 else:
     print("xformers not available")
 
-# def unload_rewriter():
-#     """Clear enhancement model from memory"""
-#     global rewriter_tokenizer, rewriter_model
-#     if rewriter_model:
-#         del rewriter_tokenizer, rewriter_model
-#         rewriter_tokenizer = None
-#         rewriter_model = None
-#     torch.cuda.empty_cache()
-#     gc.collect()
+
 @spaces.GPU()
 def infer(
     image,
@@ -273,33 +247,28 @@ def infer(
     num_inference_steps=8,
     rewrite_prompt=True,
     num_images_per_prompt=1,
+    preset_type=None,  # New parameter for presets
     progress=gr.Progress(track_tqdm=True),
 ):
     """Image editing endpoint with optimized prompt handling"""
-
     # 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"""
         try:
             if pil_image is None:
                 return pil_image
-                
             width, height = pil_image.size
             max_dimension = max(width, height)
-            
             if max_dimension <= max_size:
                 return pil_image  # No resize needed
-                
             # Calculate new dimensions maintaining aspect ratio
             scale = max_size / max_dimension
             new_width = int(width * scale)
             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}")
             return resized_image
-            
         except Exception as e:
             print(f"⚠️ Image resize failed: {e}")
             return pil_image  # Return original if resize fails
@@ -310,7 +279,6 @@ def infer(
         try:
             if pil_image is None:
                 return pil_image
-                
             img_array = np.array(pil_image).astype(np.float32) / 255.0
             noise = np.random.normal(0, noise_level, img_array.shape)
             noisy_array = img_array + noise
@@ -322,96 +290,105 @@ def infer(
         except Exception as e:
             print(f"Warning: Could not add noise to image: {e}")
             return pil_image  # Return original if noise addition fails
-    
+
     # Resize input image first
     image = resize_image(image, max_size=1024)
-    
     original_prompt = prompt
     prompt_info = ""
     
-    # Handle prompt rewriting
-    if rewrite_prompt:
-        try:
-            enhanced_instruction = polish_prompt(original_prompt)
-            if enhanced_instruction and enhanced_instruction != original_prompt:
-                prompt_info = (
-                    f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #4CAF50; background: #f5f9fe'>"
-                    f"<h4 style='margin-top: 0;'>🚀 Prompt Enhancement</h4>"
-                    f"<p><strong>Original:</strong> {original_prompt}</p>"
-                    f"<p><strong style='color:#2E7D32;'>Enhanced:</strong> {enhanced_instruction}</p>"
-                    f"</div>"
-                )
-                prompt = enhanced_instruction
-            else:
+    # Handle preset batch generation
+    if preset_type and preset_type in PRESETS:
+        preset = PRESETS[preset_type]
+        batch_prompts = [f"{original_prompt}, {preset_prompt}" for preset_prompt in preset["prompts"]]
+        num_images_per_prompt = preset["count"]
+        prompt_info = (
+            f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #2196F3; background: #f0f8ff'>"
+            f"<h4 style='margin-top: 0;'>🎨 Preset: {preset_type}</h4>"
+            f"<p>{preset['description']}</p>"
+            f"<p><strong>Base Prompt:</strong> {original_prompt}</p>"
+            f"</div>"
+        )
+        print(f"Using preset: {preset_type} with {len(batch_prompts)} variations")
+    else:
+        batch_prompts = [prompt]  # Single prompt in list
+        
+        # Handle regular prompt rewriting
+        if rewrite_prompt:
+            try:
+                enhanced_instruction = polish_prompt(original_prompt)
+                if enhanced_instruction and enhanced_instruction != original_prompt:
+                    prompt_info = (
+                        f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #4CAF50; background: #f5f9fe'>"
+                        f"<h4 style='margin-top: 0;'>🚀 Prompt Enhancement</h4>"
+                        f"<p><strong>Original:</strong> {original_prompt}</p>"
+                        f"<p><strong style='color:#2E7D32;'>Enhanced:</strong> {enhanced_instruction}</p>"
+                        f"</div>"
+                    )
+                    batch_prompts = [enhanced_instruction]
+                else:
+                    prompt_info = (
+                        f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #FF9800; background: #fff8f0'>"
+                        f"<h4 style='margin-top: 0;'>📝 Prompt Enhancement</h4>"
+                        f"<p>No enhancement applied or enhancement failed</p>"
+                        f"</div>"
+                    )
+            except Exception as e:
+                print(f"Prompt enhancement error: {str(e)}")  # Debug logging
+                gr.Warning(f"Prompt enhancement failed: {str(e)}")
                 prompt_info = (
-                    f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #FF9800; background: #fff8f0'>"
-                    f"<h4 style='margin-top: 0;'>📝 Prompt Enhancement</h4>"
-                    f"<p>No enhancement applied or enhancement failed</p>"
+                    f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #FF5252; background: #fef5f5'>"
+                    f"<h4 style='margin-top: 0;'>⚠️ Enhancement Not Applied</h4>"
+                    f"<p>Using original prompt. Error: {str(e)[:100]}</p>"
                     f"</div>"
                 )
-        except Exception as e:
-            print(f"Prompt enhancement error: {str(e)}")  # Debug logging
-            gr.Warning(f"Prompt enhancement failed: {str(e)}")
+        else:
             prompt_info = (
-                f"<div style='margin:10px; padding:15px; border-radius:8px; border-left:4px solid #FF5252; background: #fef5f5'>"
-                f"<h4 style='margin-top: 0;'>⚠️ Enhancement Not Applied</h4>"
-                f"<p>Using original prompt. Error: {str(e)[:100]}</p>"
+                f"<div style='margin:10px; padding:10px; border-radius:8px; background: #f8f9fa'>"
+                f"<h4 style='margin-top: 0;'>📝 Original Prompt</h4>"
+                f"<p>{original_prompt}</p>"
                 f"</div>"
             )
-    else:
-        prompt_info = (
-            f"<div style='margin:10px; padding:10px; border-radius:8px; background: #f8f9fa'>"
-            f"<h4 style='margin-top: 0;'>📝 Original Prompt</h4>"
-            f"<p>{original_prompt}</p>"
-            f"</div>"
-        )
     
     # Set base seed for reproducibility
     base_seed = seed if not randomize_seed else random.randint(0, MAX_SEED)
     
     try:
-        # Generate images with variation for batch mode
-        if num_images_per_prompt > 1:
-            edited_images = []
-            for i in range(num_images_per_prompt):
-                # 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
-                noisy_image = add_noise_to_image(image, noise_level=0.05 + i*0.003)
-                
-                # Slightly vary guidance scale
-                varied_guidance = true_guidance_scale + random.uniform(-0.5, 0.5)
-                varied_guidance = max(1.0, min(10.0, varied_guidance))
-                
-                # Generate single image with variations
-                result = pipe(
-                    image=noisy_image,
-                    prompt=prompt,
-                    negative_prompt=" ",
-                    num_inference_steps=num_inference_steps,
-                    generator=generator,
-                    true_cfg_scale=varied_guidance,
-                    num_images_per_prompt=1
-                ).images
-                edited_images.extend(result)
-        else:
-            # Single image generation (unchanged)
-            generator = torch.Generator(device=device).manual_seed(base_seed)
-            edited_images = pipe(
-                image=image,
-                prompt=prompt,
+        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:
+                input_image = image
+            else:
+                input_image = add_noise_to_image(image, noise_level=0.01 + i*0.003)
+            
+            # Slightly vary guidance scale for each image
+            varied_guidance = true_guidance_scale + random.uniform(-0.2, 0.2)
+            varied_guidance = max(1.0, min(10.0, varied_guidance))
+            
+            # Generate single image
+            result = pipe(
+                image=input_image,
+                prompt=current_prompt,
                 negative_prompt=" ",
                 num_inference_steps=num_inference_steps,
                 generator=generator,
-                true_cfg_scale=true_guidance_scale,
-                num_images_per_prompt=num_images_per_prompt
+                true_cfg_scale=varied_guidance,
+                num_images_per_prompt=1
             ).images
+            edited_images.extend(result)
+            
+            print(f"Generated image {i+1}/{len(batch_prompts)} with prompt: {current_prompt[:50]}...")
         
         # 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
@@ -425,13 +402,14 @@ def infer(
             f"<p>{str(e)[:200]}</p>"
             f"</div>"
         )
-        
+
+
 with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
     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>
         <p>✨ 8-step inferencing with lightx2v's LoRA.</p>
-        <p>📝 Local Prompt Enhancement, Batched Multi-image Generation</p>
+        <p>📝 Local Prompt Enhancement, Batched Multi-image Generation, 🎨 Preset Batches</p>
     </div>
     """)
     
@@ -439,65 +417,72 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
         # Input Column
         with gr.Column(scale=1):
             input_image = gr.Image(
-                label="Source Image", 
-                type="pil", 
+                label="Source Image",
+                type="pil",
                 height=300
             )
             prompt = gr.Textbox(
-                label="Edit Instructions", 
+                label="Edit Instructions",
                 placeholder="e.g. Replace the background with a beach sunset...",
                 lines=2,
                 max_lines=4
             )
             
+            # Add preset dropdown
             with gr.Row():
+                preset_dropdown = gr.Dropdown(
+                    choices=get_preset_choices(),
+                    value=None,
+                    label="Preset Batch Generation",
+                    interactive=True
+                )
                 rewrite_toggle = gr.Checkbox(
-                    label="Enable Prompt Enhancement", 
+                    label="Enable Prompt Enhancement",
                     value=True,
                     interactive=True
                 )
                 run_button = gr.Button(
-                    "Generate Edits", 
-                    variant="primary", 
+                    "Generate Edits",
+                    variant="primary",
                     min_width=120
                 )
             
             with gr.Accordion("Advanced Parameters", open=False):
                 with gr.Row():
                     seed = gr.Slider(
-                        label="Seed", 
-                        minimum=0, 
-                        maximum=MAX_SEED, 
-                        step=1, 
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
                         value=42
                     )
                     randomize_seed = gr.Checkbox(
-                        label="Random Seed", 
+                        label="Random Seed",
                         value=True
                     )
                 with gr.Row():
                     true_guidance_scale = gr.Slider(
-                        label="Guidance Scale", 
-                        minimum=1.0, 
-                        maximum=10.0, 
-                        step=0.1, 
+                        label="Guidance Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.1,
                         value=4.0
                     )
                     num_inference_steps = gr.Slider(
-                        label="Inference Steps", 
-                        minimum=4, 
-                        maximum=16, 
-                        step=1, 
+                        label="Inference Steps",
+                        minimum=4,
+                        maximum=16,
+                        step=1,
                         value=8
                     )
                 num_images_per_prompt = gr.Slider(
-                    label="Output Count", 
-                    minimum=1, 
-                    maximum=4, 
-                    step=1, 
+                    label="Output Count (Manual)",
+                    minimum=1,
+                    maximum=4,
+                    step=1,
                     value=2
                 )
-            
+                
         # Output Column
         with gr.Column(scale=2):
             result = gr.Gallery(
@@ -512,18 +497,6 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
                 "Prompt details will appear after generation</div>"
             )
     
-    # # Examples
-    # gr.Examples(
-    #     examples=[
-    #         "Change the background scene to a rooftop bar at night",
-    #         "Transform to pixel art style with 8-bit graphics",
-    #         "Replace all text with 'Qwen AI' in futuristic font"
-    #     ],
-    #     inputs=[prompt],
-    #     label="Sample Instructions",
-    #     cache_examples=True
-    # )
-
     # Set up processing
     inputs = [
         input_image,
@@ -533,9 +506,9 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
         true_guidance_scale,
         num_inference_steps,
         rewrite_toggle,
-        num_images_per_prompt
+        num_images_per_prompt,
+        preset_dropdown  # Add preset dropdown to inputs
     ]
-    
     outputs = [result, seed, prompt_info]
     
     run_button.click(
@@ -543,11 +516,11 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
         inputs=inputs,
         outputs=outputs
     )
-    
     prompt.submit(
         fn=infer,
         inputs=inputs,
         outputs=outputs
     )
 
+
 demo.queue(max_size=5).launch()