Update app_local.py

app_local.py

@@ -18,11 +18,13 @@ os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
 
 # Model configuration
 REWRITER_MODEL = "Qwen/Qwen1.5-7B-Chat"  # Upgraded to 7B for better JSON handling
-rewriter_tokenizer = None
-rewriter_model = None
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Preload enhancement model at startup
+print("🔄 Loading prompt enhancement model...")
+rewriter_tokenizer = AutoTokenizer.from_pretrained(REWRITER_MODEL)
+
 # Quantization configuration
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -31,19 +33,14 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True
 )
 
-def load_rewriter():
-    """Lazily load the prompt enhancement model"""
-    global rewriter_tokenizer, rewriter_model
-    if rewriter_tokenizer is None or rewriter_model is None:
-        print("🔄 Loading enhancement model...")
-        rewriter_tokenizer = AutoTokenizer.from_pretrained(REWRITER_MODEL)
-        rewriter_model = AutoModelForCausalLM.from_pretrained(
-            REWRITER_MODEL,
-            torch_dtype=dtype,
-            device_map="auto",
-            quantization_config=bnb_config
-        )
-        print("✅ Enhancement model loaded")
+rewriter_model = AutoModelForCausalLM.from_pretrained(
+    REWRITER_MODEL,
+    torch_dtype=dtype,
+    device_map="auto",
+    quantization_config=bnb_config,
+    max_memory={0: "48GiB"},  # Reserve adequate memory
+)
+print("✅ Enhancement model loaded and ready!")
 
 SYSTEM_PROMPT_EDIT = '''
 # Edit Instruction Rewriter
@@ -103,7 +100,7 @@ def extract_json_response(model_output: str) -> str:
         
         # Extract rewritten prompt from possible key variations
         possible_keys = [
-            "Rewritten", "rewritten", "Rewrited", "rewrited",
+            "Rewritten", "rewritten", "Rewrited", "rewrited", "Rewrittent",
             "Output", "output", "Enhanced", "enhanced"
         ]
         for key in possible_keys:
@@ -151,9 +148,9 @@ def polish_prompt(original_prompt: str) -> str:
     with torch.no_grad():
         generated_ids = rewriter_model.generate(
             **model_inputs,
-            max_new_tokens=150,  # Reduced for better quality
+            max_new_tokens=256,  # Reduced for better quality
             do_sample=True,
-            temperature=0.4,  # Less creative but more focused
+            temperature=0.5,  # Less creative but more focused
             top_p=0.9,
             no_repeat_ngram_size=3,
             pad_token_id=rewriter_tokenizer.eos_token_id
@@ -218,15 +215,15 @@ 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()
+# 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(duration=60)
 def infer(
@@ -271,9 +268,7 @@ def infer(
             f"</div>"
         )
     
-    # Free VRAM after enhancement
-    unload_rewriter()
-    
+
     # Set seed for reproducibility
     seed_val = seed if not randomize_seed else random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed_val)
@@ -302,11 +297,12 @@ def infer(
 
 MAX_SEED = np.iinfo(np.int32).max
 
-with gr.Blocks(title="Qwen Image Editor Fast", css=".gr-gallery {min-height: 300px}") as demo:
+with gr.Blocks(title="Qwen Image Editor Fast") 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 • Local Prompt Enhancement • H200 Optimized</p>
+        <p>✨ 8-step inferencing with lightx2v's LoRA.")
+        <p>📝 Local Prompt Enhancement</p>
     </div>
     """)