Update app_local.py

app_local.py

@@ -170,6 +170,7 @@ 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
     # Try to extract JSON content
     rewritten_prompt = extract_json_response(enhanced)
@@ -227,15 +228,20 @@ pipe = QwenImageEditPipeline.from_pretrained(
 # Load LoRA weights for acceleration
 pipe.load_lora_weights(
     "lightx2v/Qwen-Image-Lightning",
-    weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
+    # weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
+    weight_name="Qwen-Image-Lightning-4steps-V1.0.safetensors"
 )
 pipe.fuse_lora()
 
-if is_xformers_available():
-    pipe.enable_xformers_memory_efficient_attention()
-else:
-    print("xformers not available")
-
+# if is_xformers_available():
+#     pipe.enable_xformers_memory_efficient_attention()
+# else:
+#     print("xformers not available")
+try:
+    pipe.enable_vae_slicing()
+except Exception as e:
+    print(f"VAE Slicing Failed: {e}")
+    
 # Function to update prompt preview when preset is selected
 def update_prompt_preview(preset_type, base_prompt):
     """Update the prompt preview display based on selected preset and base prompt"""
@@ -261,7 +267,7 @@ def infer(
     seed=42,
     randomize_seed=False,
     true_guidance_scale=4.0,
-    num_inference_steps=8,
+    num_inference_steps=4,
     rewrite_prompt=True,
     num_images_per_prompt=1,
     preset_type=None,  # New parameter for presets
@@ -291,7 +297,7 @@ def infer(
             return pil_image  # Return original if resize fails
     
     # Add noise function for batch variation
-    def add_noise_to_image(pil_image, noise_level=0.05):
+    def add_noise_to_image(pil_image, noise_level=0.02):
         """Add slight noise to image to create variation in outputs"""
         try:
             if pil_image is None:
@@ -384,7 +390,7 @@ def infer(
                 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 = true_guidance_scale + random.uniform(-0.1, 0.1)
             varied_guidance = max(1.0, min(10.0, varied_guidance))
             
             # Generate single image
@@ -399,12 +405,12 @@ def infer(
             ).images
             edited_images.extend(result)
             
-            print(f"Generated image {i+1}/{len(batch_prompts)} with prompt: {current_prompt[:50]}...")
+            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()
+        # if device == "cuda":
+        #     torch.cuda.empty_cache()
+        #     gc.collect()
             
         return edited_images, base_seed, prompt_info
     except Exception as e:
@@ -487,7 +493,7 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
                     )
                 with gr.Row():
                     true_guidance_scale = gr.Slider(
-                        label="Guidance Scale",
+                        label="True CFG Scale",
                         minimum=1.0,
                         maximum=10.0,
                         step=0.1,
@@ -495,10 +501,10 @@ with gr.Blocks(title="Qwen Image Edit - Fast Lightning Mode w/ Batch") as demo:
                     )
                     num_inference_steps = gr.Slider(
                         label="Inference Steps",
-                        minimum=4,
+                        minimum=2,
                         maximum=16,
                         step=1,
-                        value=6
+                        value=4
                     )
                 num_images_per_prompt = gr.Slider(
                     label="Output Count (Manual)",