Update app.py

app.py

@@ -5,428 +5,155 @@ from model import ModelHandler
 from generator import Generator
 from config import Config
 
-# ============================================================================
-# INITIALIZATION
-# ============================================================================
-
-print("🚀 Initializing Enhanced Face To Style Application...")
+# 1. Initialize Models Globally
+print("Initializing Application...")
 handler = ModelHandler()
-
-# Load with default edge type (can be changed via environment variable)
-import os
-edge_type = os.getenv("EDGE_TYPE", Config.DEFAULT_EDGE_TYPE)
-print(f"Loading models with edge type: {edge_type}")
-handler.load_models(edge_type=edge_type)
-
+handler.load_models()
 gen = Generator(handler)
-print("✅ Application Ready!\n")
-
-# ============================================================================
-# INFERENCE FUNCTION
-# ============================================================================
 
+# 2. Define GPU-enabled Inference Function
 @spaces.GPU(duration=20)
 def process_img(
     image, 
     prompt,
     negative_prompt,
-    cfg_scale,
+    cfg_scale, # <-- RE-ENABLED
     steps, 
     img_strength, 
     depth_strength, 
     edge_strength,
-    instantid_strength,
-    canny_low,
-    canny_high,
-    eta,
-    seed,
-    show_controls
+    seed
 ):
-    """Enhanced processing with control image preview option"""
     if image is None:
         raise gr.Error("Please upload an image first.")
     
     try:
-        print("=" * 70)
-        print("STARTING GENERATION")
-        print("=" * 70)
-        
+        print("--- Starting Generation ---")
         result = gen.predict(
             image, 
             prompt,
             negative_prompt=negative_prompt,
-            guidance_scale=cfg_scale,
+            guidance_scale=cfg_scale, # <-- RE-ENABLED
             num_inference_steps=steps,
             img2img_strength=img_strength,
             depth_strength=depth_strength,
-            edge_strength=edge_strength,
-            instantid_strength=instantid_strength,
-            canny_low_threshold=canny_low,
-            canny_high_threshold=canny_high,
-            eta=eta,
-            seed=seed,
-            return_control_images=show_controls
+            lineart_strength=edge_strength,
+            seed=seed
         )
-        
-        print("=" * 70)
-        print("✅ GENERATION COMPLETE")
-        print("=" * 70)
-        
-        if show_controls:
-            output_image, control_dict = result
-            # Create a preview grid of control images
-            controls_preview = create_control_preview(control_dict)
-            return output_image, controls_preview
-        else:
-            return result, None
+        print("--- Generation Complete ---")
+        return result
         
     except Exception as e:
-        print(f"❌ Error during generation: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"Error during generation: {e}")
         raise gr.Error(f"An error occurred: {str(e)}")
 
-def create_control_preview(control_dict):
-    """Create a preview grid of control images"""
-    from PIL import Image
-    import numpy as np
-    
-    images = []
-    labels = []
-    
-    if control_dict.get('processed_input'):
-        images.append(control_dict['processed_input'])
-        labels.append("Processed Input")
-    
-    if control_dict.get('depth'):
-        images.append(control_dict['depth'])
-        labels.append("Depth Map")
-    
-    if control_dict.get('edges'):
-        for i, edge in enumerate(control_dict['edges']):
-            images.append(edge)
-            edge_type = "Canny" if i == 0 else "LineArt"
-            labels.append(f"{edge_type} Edges")
-    
-    if control_dict.get('face_kps'):
-        images.append(control_dict['face_kps'])
-        labels.append("Face Keypoints")
-    
-    # Create grid
-    if not images:
-        return None
-    
-    # Simple horizontal concatenation for now
-    # You could make this fancier with a proper grid layout
-    return images[0]  # Just return first for now - can be enhanced
-
-# ============================================================================
-# PRESET CONFIGURATIONS
-# ============================================================================
-
-PRESETS = {
-    "Balanced (Default)": {
-        "cfg_scale": 4.0,
-        "steps": 8,
-        "img_strength": 0.8,
-        "depth_strength": 0.5,
-        "edge_strength": 0.5,
-        "instantid_strength": 0.8,
-        "eta": 0.45
-    },
-    "Fast": {
-        "cfg_scale": 3.0,
-        "steps": 4,
-        "img_strength": 0.7,
-        "depth_strength": 0.4,
-        "edge_strength": 0.4,
-        "instantid_strength": 0.7,
-        "eta": 0.5
-    },
-    "Quality": {
-        "cfg_scale": 5.0,
-        "steps": 12,
-        "img_strength": 0.85,
-        "depth_strength": 0.6,
-        "edge_strength": 0.6,
-        "instantid_strength": 0.85,
-        "eta": 0.3
-    },
-    "Strong Style": {
-        "cfg_scale": 7.0,
-        "steps": 10,
-        "img_strength": 0.75,
-        "depth_strength": 0.4,
-        "edge_strength": 0.4,
-        "instantid_strength": 0.6,
-        "eta": 0.4
-    },
-    "Preserve Structure": {
-        "cfg_scale": 4.0,
-        "steps": 8,
-        "img_strength": 0.9,
-        "depth_strength": 0.8,
-        "edge_strength": 0.8,
-        "instantid_strength": 0.9,
-        "eta": 0.3
-    }
-}
-
-def apply_preset(preset_name):
-    """Apply preset configuration"""
-    if preset_name in PRESETS:
-        preset = PRESETS[preset_name]
-        return [
-            preset["cfg_scale"],
-            preset["steps"],
-            preset["img_strength"],
-            preset["depth_strength"],
-            preset["edge_strength"],
-            preset["instantid_strength"],
-            preset["eta"]
-        ]
-    return [4.0, 8, 0.8, 0.5, 0.5, 0.8, 0.45]  # Default
-
-# ============================================================================
-# GRADIO INTERFACE
-# ============================================================================
-
-with gr.Blocks(title="Enhanced Face To Style", theme=gr.themes.Soft()) as demo:
+# 3. Build Gradio Interface
+with gr.Blocks(title="Face To Style", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🎨 Enhanced Face to Style
-        
-        Upload any image and transform it with AI-powered style transfer!
-        - **With Face**: Preserves identity using InstantID
-        - **Without Face**: Pure artistic stylization
-        - **TCD Accelerated**: Ultra-fast generation (4-12 steps)
+        # 🎨 Face to Style
+        Upload any image. If there is a face, we'll keep the identity. If not, we'll stylize the scene!
         
-        ---
+        **Using DPMSolver++ scheduler with Karras sigmas for high-quality results.**
         """
     )
     
     with gr.Row():
-        # ===== LEFT COLUMN: INPUTS =====
         with gr.Column(scale=2):
-            input_img = gr.Image(type="pil", label="📸 Input Image")
-            
+            input_img = gr.Image(type="pil", label="Input Image")
             prompt = gr.Textbox(
-                label="✍️ Prompt (Optional)", 
+                label="Prompt (Optional)", 
                 placeholder="Leave empty for auto-captioning...",
-                info=f"Trigger words '{Config.STYLE_TRIGGER}' are added automatically.",
-                lines=2
+                info=f"The trigger words '{Config.STYLE_TRIGGER}' are added automatically."
             )
             
             negative_prompt = gr.Textbox(
-                label="🚫 Negative Prompt", 
+                label="Negative Prompt (Optional)", 
                 placeholder="e.g., blurry, text, watermark, bad art...",
-                value=Config.DEFAULT_NEGATIVE_PROMPT,
-                lines=2
+                value=Config.DEFAULT_NEGATIVE_PROMPT
             )
             
-            # Preset Selector
-            with gr.Row():
-                preset_dropdown = gr.Dropdown(
-                    choices=list(PRESETS.keys()),
-                    value="Balanced (Default)",
-                    label="🎛️ Presets",
-                    info="Quick configurations for different use cases"
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Number(
+                    label="Seed", 
+                    value=-1, 
+                    info="-1 for random", 
+                    precision=0
                 )
-                apply_preset_btn = gr.Button("Apply Preset", size="sm")
-            
-            # Advanced Settings
-            with gr.Accordion("🔧 Advanced Settings", open=False):
-                with gr.Row():
-                    seed = gr.Number(
-                        label="🎲 Seed", 
-                        value=-1, 
-                        info="-1 for random", 
-                        precision=0
-                    )
-                    show_controls = gr.Checkbox(
-                        label="Show Control Images",
-                        value=False,
-                        info="Display depth/edge maps"
-                    )
                 
-                gr.Markdown("### Core Parameters")
-                with gr.Row():
-                    cfg_scale = gr.Slider(
-                        minimum=1.0, 
-                        maximum=10.0,
-                        step=0.1, 
-                        value=4.0,
-                        label="🎨 Style Strength (CFG)",
-                        info="Higher = stronger style adherence"
-                    )
-                    
-                    steps = gr.Slider(
-                        minimum=1, 
-                        maximum=20, 
-                        step=1, 
-                        value=8,
-                        label="⚡ Steps",
-                        info="TCD works well with 4-12"
-                    )
-                
-                with gr.Row():
-                    img_strength = gr.Slider(
-                        minimum=0.1, 
-                        maximum=1.0, 
-                        step=0.05, 
-                        value=0.8,
-                        label="🖼️ Image Strength",
-                        info="Higher = more transformation"
-                    )
-                    
-                    eta = gr.Slider(
-                        minimum=0.0,
-                        maximum=1.0,
-                        step=0.05,
-                        value=0.45,
-                        label="🎲 Eta (Randomness)",
-                        info="TCD stochasticity parameter"
-                    )
-                
-                gr.Markdown("### Control Strengths")
-                with gr.Row():
-                    depth_strength = gr.Slider(
-                        minimum=0.0, 
-                        maximum=1.0, 
-                        step=0.05, 
-                        value=0.5,
-                        label="🗺️ Depth Strength",
-                        info="Spatial structure preservation"
-                    )
-                    
-                    edge_strength = gr.Slider(
-                        minimum=0.0, 
-                        maximum=1.0, 
-                        step=0.05, 
-                        value=0.5,
-                        label="✏️ Edge Strength",
-                        info="Boundary/detail preservation"
-                    )
-                
-                instantid_strength = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.05,
-                    value=0.8,
-                    label="👤 Face Preservation (InstantID)",
-                    info="Only applies if face detected"
+                # --- CFG/GUIDANCE SLIDER ---
+                cfg_scale = gr.Slider(
+                    elem_id="cfg_scale",
+                    minimum=1.0, 
+                    maximum=15.0, 
+                    step=0.5, 
+                    value=Config.CGF_SCALE, # Default 7.0
+                    label="Guidance Scale (CFG)"
                 )
                 
-                # Canny-specific controls (only show if using canny)
-                if edge_type in ["canny", "both"]:
-                    gr.Markdown("### Canny Edge Detection")
-                    with gr.Row():
-                        canny_low = gr.Slider(
-                            minimum=0,
-                            maximum=255,
-                            step=1,
-                            value=100,
-                            label="Low Threshold",
-                            info="Lower = more edges"
-                        )
-                        canny_high = gr.Slider(
-                            minimum=0,
-                            maximum=255,
-                            step=1,
-                            value=200,
-                            label="High Threshold",
-                            info="Higher = only strong edges"
-                        )
-                else:
-                    # Hidden sliders with default values
-                    canny_low = gr.Slider(value=100, visible=False)
-                    canny_high = gr.Slider(value=200, visible=False)
+                steps = gr.Slider(
+                    elem_id="steps",
+                    minimum=15, 
+                    maximum=50, 
+                    step=1, 
+                    value=20, # DPMSolver++ default
+                    label="Steps Number"
+                )
+                img_strength = gr.Slider(
+                    elem_id="img_strength",
+                    minimum=0.5, 
+                    maximum=1.0, 
+                    step=0.05, 
+                    value=Config.IMG_STRENGTH, 
+                    label="Image Strength (Img2Img)"
+                )
+                depth_strength = gr.Slider(
+                    elem_id="depth_strength",
+                    minimum=0.0, 
+                    maximum=1.0, 
+                    step=0.05, 
+                    value=Config.DEPTH_STRENGTH, 
+                    label="DepthMap Strength"
+                )
+                edge_strength = gr.Slider(
+                    elem_id="edge_strength",
+                    minimum=0.0, 
+                    maximum=1.0, 
+                    step=0.05, 
+                    value=Config.EDGE_STRENGTH, 
+                    label="EdgeMap Strength (LineArt)"
+                )
             
-            run_btn = gr.Button("✨ Generate", variant="primary", size="lg")
+            run_btn = gr.Button("Generate", variant="primary")
             
-        # ===== RIGHT COLUMN: OUTPUTS =====
         with gr.Column(scale=1):
-            output_img = gr.Image(label="🎨 Styled Result")
-            control_preview = gr.Image(label="🔍 Control Images Preview", visible=False)
-    
-    # Info Section
-    with gr.Row():
-        gr.Markdown(
-            f"""
-            ### 💡 Tips
-            - **Fast Preview**: Use 4-6 steps with Fast preset
-            - **Best Quality**: Use 10-12 steps with Quality preset
-            - **Strong Stylization**: Increase CFG scale and decrease control strengths
-            - **Preserve Details**: Increase depth/edge strengths
-            - **Edge Type**: Currently using **{edge_type.upper()}** detection
-            
-            ### ⚙️ Current Configuration
-            - **LoRA**: {Config.LORA_FILENAME} (strength: {Config.LORA_STRENGTH})
-            - **Trigger**: `{Config.STYLE_TRIGGER}`
-            - **Scheduler**: TCD (Trajectory Consistency Distillation)
-            """
-        )
-    
-    # ===== EVENT HANDLERS =====
-    
-    # Collect all inputs
+            output_img = gr.Image(label="Styled Result")
+                
+    # Event Handler
     all_inputs = [
         input_img, 
         prompt, 
         negative_prompt,
-        cfg_scale,
+        cfg_scale, # <-- RE-ENABLED
         steps, 
         img_strength, 
         depth_strength, 
         edge_strength,
-        instantid_strength,
-        canny_low,
-        canny_high,
-        eta,
-        seed,
-        show_controls
+        seed
     ]
     
-    # Generate button
     run_btn.click(
         fn=process_img, 
         inputs=all_inputs, 
-        outputs=[output_img, control_preview]
-    )
-    
-    # Preset application
-    preset_outputs = [
-        cfg_scale,
-        steps,
-        img_strength,
-        depth_strength,
-        edge_strength,
-        instantid_strength,
-        eta
-    ]
-    
-    apply_preset_btn.click(
-        fn=apply_preset,
-        inputs=[preset_dropdown],
-        outputs=preset_outputs
-    )
-    
-    # Show/hide control preview based on checkbox
-    show_controls.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=[show_controls],
-        outputs=[control_preview]
+        outputs=[output_img]
     )
 
-# ============================================================================
-# LAUNCH
-# ============================================================================
 
+# 4. Launch the App
 if __name__ == "__main__":
     demo.queue(max_size=20, api_open=True)
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        show_api=True,
-        share=False
+        show_api=True 
     )