Update app_local.py

app_local.py

@@ -16,7 +16,7 @@ os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
 os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
 
 # Model configuration
-REWRITER_MODEL = "Qwen/Qwen1.5-1.8B-Chat"
+REWRITER_MODEL = "Qwen/Qwen1.5-7B-Chat"  # Upgraded to 7B for better JSON handling
 rewriter_tokenizer = None
 rewriter_model = None
 dtype = torch.bfloat16
@@ -76,17 +76,55 @@ 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"""
+    try:
+        # Try to find the JSON portion in the output
+        start_idx = model_output.find('{')
+        end_idx = model_output.rfind('}') + 1
+        if start_idx == -1 or end_idx == 0:
+            return None
+        
+        json_str = model_output[start_idx:end_idx]
+        # Clean up common formatting issues
+        json_str = re.sub(r'(?<!")\b(\w+)\b(?=":)', r'"\1"', json_str)  # Add quotes to keys
+        json_str = re.sub(r':\s*([^"{\[]|true|false|null)', r': "\1"', json_str)  # Add quotes to values
+        
+        # Parse JSON
+        data = json.loads(json_str)
+        
+        # Extract rewritten prompt from possible key variations
+        possible_keys = [
+            "Rewritten", "rewritten", "Rewrited", "rewrited", 
+            "Output", "output", "Enhanced", "enhanced"
+        ]
+        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()
+            
+        # Fallback to direct extraction
+        for value in data.values():
+            if isinstance(value, str) and 10 < len(value) < 500:
+                return value.strip()
+                
+    except Exception:
+        pass
+    return None
+    
 def polish_prompt(original_prompt: str) -> str:
-    """Enhanced prompt rewriting using Qwen1.5-1.8B"""
+    """Enhanced prompt rewriting using original system prompt with JSON handling"""
     load_rewriter()
     
-    # Format as Qwen chat with system prompt
+    # Format as Qwen chat
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT_EDIT},
         {"role": "user", "content": original_prompt}
     ]
     
-    # Generate enhanced prompt
     text = rewriter_tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -98,29 +136,40 @@ def polish_prompt(original_prompt: str) -> str:
     with torch.no_grad():
         generated_ids = rewriter_model.generate(
             **model_inputs,
-            max_new_tokens=120,
+            max_new_tokens=256,  # Maintain token count for good JSON generation
             do_sample=True,
-            temperature=0.7,
-            top_p=0.95,
-            no_repeat_ngram_size=2
+            temperature=0.6,
+            top_p=0.9,
+            no_repeat_ngram_size=2,
+            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()
     
-    # Clean possible artifacts
-    enhanced = enhanced.strip()
-    if enhanced.lower().startswith(("rewritten instruction:", "enhanced:", "output:")):
-        enhanced = re.split(r':', enhanced, 1)[-1].strip()
+    # Try to extract JSON content
+    rewritten_prompt = extract_json_response(enhanced)
     
-    # Remove any quotes around the prompt if present
-    if enhanced.startswith('"') and enhanced.endswith('"'):
-        enhanced = enhanced[1:-1]
+    if rewritten_prompt:
+        # Clean up substitutions from the JSON output
+        rewritten_prompt = re.sub(r'(Replace|Change|Add) "([^"]*)"', r'\1 \2', rewritten_prompt)
+        rewritten_prompt = rewritten_prompt.replace('\\"', '"')
+        return rewritten_prompt
     
-    return enhanced
+    # Fallback cleanup if JSON extraction fails
+    print(f"⚠️ JSON extraction failed, using raw output: {enhanced}")
+    fallback = re.sub(r'```.*?```', '', enhanced, flags=re.DOTALL)  # Remove code blocks
+    fallback = re.sub(r'[\{\}\[\]"]', '', fallback)  # Remove JSON artifacts
+    fallback = fallback.split('\n')[0]  # Take first line
+    
+    # Try to extract before colon separator
+    if ': ' in fallback:
+        return fallback.split(': ')[1].strip()
+    
+    return fallback.strip()
 
 # Load main image editing pipeline
 pipe = QwenImageEditPipeline.from_pretrained(
@@ -140,6 +189,7 @@ if is_xformers_available():
 else:
     print("xformers not available")
 
+
 def unload_rewriter():
     """Clear enhancement model from memory"""
     global rewriter_tokenizer, rewriter_model
@@ -266,8 +316,7 @@ with gr.Blocks(title="Qwen Image Editor Fast") as demo:
             
             rewrite_toggle = gr.Checkbox(
                 label="Enable AI Prompt Enhancement", 
-                value=True,
-                info="Uses local Qwen1.5-1.8B model to improve your instructions"
+                value=True
             )
             
             run_button = gr.Button("Generate Edits", variant="primary")