Initial commit: Architecture AI Enhancer v1.0.0

README.md

@@ -0,0 +1,59 @@
+---
+title: Architecture AI Enhancer
+emoji: 🏗️
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: mit
+python_version: 3.10
+---
+
+# Architecture AI Enhancer
+
+Transform your architectural renders with AI-powered enhancement using Stable Diffusion 1.5.
+
+## Features
+
+- 🎨 AI image enhancement with Stable Diffusion
+- 📈 2x upscaling with ESRGAN
+- 🎭 Photographic post-processing
+- ⚙️ Configurable parameters
+- 🚀 GPU-accelerated inference
+
+## Usage
+
+1. Upload an architectural render
+2. Adjust enhancement settings (optional)
+3. Click "Enhance Image"
+4. Wait 30-60 seconds for processing
+5. Download your enhanced result
+
+## Technical Details
+
+- **Model**: runwayml/stable-diffusion-v1-5
+- **Framework**: PyTorch + Diffusers
+- **Interface**: Gradio 4.44.0
+- **Version**: 1.0.0
+
+## Settings Guide
+
+- **Strength** (0.1-0.8): Controls how much the AI modifies the image
+  - 0.2-0.3: Subtle enhancements
+  - 0.4-0.5: Moderate changes
+  - 0.6-0.8: Dramatic transformations
+
+- **Guidance Scale** (1-15): How closely the AI follows the prompt
+  - 4-6: Natural, balanced results
+  - 7-10: More stylized
+  - 11-15: Very strong prompt adherence
+
+## Credits
+
+Built with:
+- [Stable Diffusion](https://github.com/CompVis/stable-diffusion)
+- [Diffusers](https://github.com/huggingface/diffusers)
+- [Gradio](https://gradio.app)
+- [PyTorch](https://pytorch.org)

app.py

@@ -0,0 +1,230 @@
+"""
+Architecture AI Enhancer - Gradio App for Hugging Face Spaces
+
+This file creates a Gradio interface for the enhancement pipeline.
+"""
+import gradio as gr
+import torch
+from PIL import Image
+import logging
+from pathlib import Path
+import sys
+
+# Add backend to path
+sys.path.insert(0, str(Path(__file__).parent / "backend"))
+
+from services.diffusion_pipeline import DiffusionPipeline
+from services.upscaler import ESRGANUpscaler
+from services.post_processor import PostProcessor
+from config import settings
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize pipelines (lazy loading)
+diffusion_pipeline = None
+upscaler = None
+post_processor = None
+
+def initialize_models():
+    """Initialize all models on first use"""
+    global diffusion_pipeline, upscaler, post_processor
+    
+    if diffusion_pipeline is None:
+        logger.info("Loading Stable Diffusion pipeline...")
+        diffusion_pipeline = DiffusionPipeline()
+        
+    if upscaler is None:
+        logger.info("Loading upscaler...")
+        upscaler = ESRGANUpscaler()
+        
+    if post_processor is None:
+        post_processor = PostProcessor()
+
+def enhance_image(
+    input_image: Image.Image,
+    strength: float = 0.3,
+    guidance_scale: float = 5.5,
+    custom_prompt: str = "",
+    use_upscaler: bool = True,
+    use_postprocess: bool = True,
+    progress=gr.Progress()
+) -> Image.Image:
+    """
+    Enhance architectural image using AI pipeline
+    
+    Args:
+        input_image: Input PIL Image
+        strength: Denoising strength (0.1-0.8)
+        guidance_scale: CFG scale (1.0-15.0)
+        custom_prompt: Optional custom prompt
+        use_upscaler: Apply upscaling
+        use_postprocess: Apply post-processing
+        progress: Gradio progress tracker
+        
+    Returns:
+        Enhanced PIL Image
+    """
+    try:
+        # Initialize models
+        progress(0, desc="Loading models...")
+        initialize_models()
+        
+        # Step 1: AI Enhancement
+        progress(0.1, desc="Preprocessing image...")
+        
+        prompt = custom_prompt if custom_prompt else settings.DEFAULT_PROMPT
+        
+        progress(0.2, desc="Enhancing with AI (this may take a few minutes)...")
+        
+        enhanced = diffusion_pipeline.run_inference(
+            image=input_image,
+            prompt=prompt,
+            negative_prompt=settings.NEGATIVE_PROMPT,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=30
+        )
+        
+        progress(0.7, desc="AI enhancement complete!")
+        
+        # Step 2: Optional Upscaling
+        if use_upscaler:
+            progress(0.75, desc="Upscaling image...")
+            enhanced = upscaler.upscale(enhanced)
+            progress(0.85, desc="Upscaling complete!")
+        
+        # Step 3: Optional Post-processing
+        if use_postprocess:
+            progress(0.9, desc="Applying final touches...")
+            enhanced = post_processor.process(enhanced)
+            progress(0.95, desc="Post-processing complete!")
+        
+        progress(1.0, desc="Done!")
+        return enhanced
+        
+    except Exception as e:
+        logger.error(f"Enhancement failed: {e}", exc_info=True)
+        raise gr.Error(f"Enhancement failed: {str(e)}")
+
+# Create Gradio interface
+with gr.Blocks(title="Architecture AI Enhancer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🏗️ Architecture AI Enhancer
+    
+    Transform your architectural renders with AI-powered enhancement using Stable Diffusion 1.5.
+    
+    **Upload an image** and adjust the settings below to enhance your architectural visualization.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="📤 Input Image",
+                type="pil",
+                height=400
+            )
+            
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                strength = gr.Slider(
+                    minimum=0.1,
+                    maximum=0.8,
+                    value=0.3,
+                    step=0.05,
+                    label="Denoising Strength",
+                    info="Lower = more faithful to input, Higher = more creative"
+                )
+                
+                guidance_scale = gr.Slider(
+                    minimum=1.0,
+                    maximum=15.0,
+                    value=5.5,
+                    step=0.5,
+                    label="Guidance Scale",
+                    info="How closely to follow the prompt"
+                )
+                
+                custom_prompt = gr.Textbox(
+                    label="Custom Prompt (optional)",
+                    placeholder="professional architectural photography, detailed, high quality...",
+                    lines=3
+                )
+                
+                use_upscaler = gr.Checkbox(
+                    label="Enable Upscaling (2x)",
+                    value=True
+                )
+                
+                use_postprocess = gr.Checkbox(
+                    label="Enable Post-Processing",
+                    value=True,
+                    info="Adds photographic enhancements"
+                )
+            
+            enhance_btn = gr.Button("✨ Enhance Image", variant="primary", size="lg")
+            
+        with gr.Column():
+            output_image = gr.Image(
+                label="✅ Enhanced Result",
+                type="pil",
+                height=400
+            )
+            
+            gr.Markdown("""
+            ### 📝 Tips for Best Results:
+            - Use high-quality architectural renders as input
+            - Start with default settings and adjust if needed
+            - Lower strength for subtle enhancements
+            - Higher strength for more dramatic changes
+            - Processing takes 2-5 minutes on CPU, ~30 seconds on GPU
+            """)
+    
+    gr.Markdown("""
+    ---
+    ### 🔧 Technical Details:
+    - **Model**: Stable Diffusion 1.5
+    - **Upscaler**: ESRGAN (optional)
+    - **Processing**: CPU/GPU automatic detection
+    - **Version**: 1.0.0
+    
+    ### 📚 Resources:
+    - [GitHub Repository](#)
+    - [Documentation](#)
+    - [Report Issues](#)
+    """)
+    
+    # Connect the enhance button
+    enhance_btn.click(
+        fn=enhance_image,
+        inputs=[
+            input_image,
+            strength,
+            guidance_scale,
+            custom_prompt,
+            use_upscaler,
+            use_postprocess
+        ],
+        outputs=output_image
+    )
+    
+    # Example images
+    gr.Examples(
+        examples=[
+            ["examples/render1.jpg", 0.3, 5.5, "", True, True],
+            ["examples/render2.jpg", 0.4, 6.0, "", True, True],
+        ],
+        inputs=[input_image, strength, guidance_scale, custom_prompt, use_upscaler, use_postprocess],
+        outputs=output_image,
+        fn=enhance_image,
+        cache_examples=False
+    )
+
+# Launch configuration
+if __name__ == "__main__":
+    demo.queue(max_size=10)  # Enable queue for multiple users
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )

backend/config.py

@@ -0,0 +1,160 @@
+"""
+Configuration module for Architecture AI Enhancer
+Centralized configuration management for the entire application
+"""
+import os
+from pathlib import Path
+from typing import Optional
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    """
+    Application settings with environment variable support
+    
+    All settings can be overridden via environment variables
+    """
+    
+    # Application Metadata
+    APP_NAME: str = "Architecture AI Enhancer"
+    APP_VERSION: str = "1.0.0"
+    DEBUG: bool = False
+    
+    # Server Configuration
+    HOST: str = "0.0.0.0"
+    PORT: int = 8000
+    WORKERS: int = 1
+    
+    # CORS Settings
+    CORS_ORIGINS: list = [
+        "http://localhost:3000",
+        "http://localhost:5173",
+        "http://127.0.0.1:3000",
+        "http://127.0.0.1:5173"
+    ]
+    
+    # Path Configuration
+    BASE_DIR: Path = Path(__file__).parent
+    MODELS_DIR: Path = BASE_DIR / "models"
+    LORA_DIR: Path = MODELS_DIR / "lora"
+    BASE_MODEL_DIR: Path = MODELS_DIR / "base"
+    DATASETS_DIR: Path = BASE_DIR / "datasets"
+    INPUT_DIR: Path = DATASETS_DIR / "input"
+    TARGET_DIR: Path = DATASETS_DIR / "target"
+    PROCESSED_DIR: Path = DATASETS_DIR / "processed"
+    OUTPUT_DIR: Path = BASE_DIR / "output"
+    ENHANCED_DIR: Path = OUTPUT_DIR / "enhanced"
+    LOGS_DIR: Path = OUTPUT_DIR / "logs"
+    
+    # Model Configuration
+    BASE_MODEL: str = "runwayml/stable-diffusion-v1-5"  # Lighter model (~4GB RAM vs ~12GB for SDXL)
+    LORA_MODEL_NAME: str = "office_style.safetensors"
+    VAE_MODEL: Optional[str] = None
+    
+    # Training Hyperparameters
+    LORA_RANK: int = 8
+    LEARNING_RATE: float = 1e-4
+    TRAIN_STEPS: int = 1000
+    BATCH_SIZE: int = 1
+    GRADIENT_ACCUMULATION_STEPS: int = 4
+    MAX_GRAD_NORM: float = 1.0
+    WARMUP_STEPS: int = 100
+    SAVE_STEPS: int = 250
+    
+    # Inference Configuration
+    IMG2IMG_STRENGTH: float = 0.3  # Range: 0.2-0.35
+    GUIDANCE_SCALE: float = 5.5    # Range: 4-7
+    NUM_INFERENCE_STEPS: int = 30
+    
+    # Image Processing
+    MAX_IMAGE_SIZE: int = 2048
+    UPSCALE_FACTOR: int = 2
+    TARGET_RESOLUTION: int = 512  # SD 1.5 works best at 512x512
+    
+    # Prompts
+    DEFAULT_PROMPT: str = (
+        "ultra realistic architectural visualization, "
+        "professional photography, high detail, sharp focus, "
+        "natural lighting, modern office interior, "
+        "clean lines, photorealistic rendering"
+    )
+    
+    NEGATIVE_PROMPT: str = (
+        "distorted walls, cartoon, illustration, painting, drawing, "
+        "unrealistic proportions, blurry, low quality, artifacts, "
+        "oversaturated, noise, grain, ugly, deformed"
+    )
+    
+    # Device Configuration
+    DEVICE: str = "cuda"  # "cuda" or "cpu" - auto-detected at runtime
+    MIXED_PRECISION: str = "fp16"  # "fp16", "bf16", or "no"
+    ENABLE_ATTENTION_SLICING: bool = True  # Reduce VRAM usage
+    ENABLE_VAE_SLICING: bool = True  # Reduce VRAM usage
+    
+    # Upload Limits
+    MAX_UPLOAD_SIZE: int = 25 * 1024 * 1024  # 25 MB
+    ALLOWED_EXTENSIONS: set = {".png", ".jpg", ".jpeg", ".webp"}
+    
+    class Config:
+        env_file = ".env"
+        case_sensitive = True
+
+
+# Global settings instance
+settings = Settings()
+
+
+def ensure_directories():
+    """
+    Create all necessary directories if they don't exist
+    
+    This function should be called on application startup
+    """
+    directories = [
+        settings.MODELS_DIR,
+        settings.LORA_DIR,
+        settings.BASE_MODEL_DIR,
+        settings.DATASETS_DIR,
+        settings.INPUT_DIR,
+        settings.TARGET_DIR,
+        settings.PROCESSED_DIR,
+        settings.OUTPUT_DIR,
+        settings.ENHANCED_DIR,
+        settings.LOGS_DIR,
+    ]
+    
+    for directory in directories:
+        directory.mkdir(parents=True, exist_ok=True)
+        print(f"✓ Ensured directory exists: {directory}")
+
+
+def get_lora_path() -> Optional[Path]:
+    """
+    Get the path to the trained LoRA model if it exists
+    
+    Returns:
+        Path to LoRA model or None if not found
+    """
+    lora_path = settings.LORA_DIR / settings.LORA_MODEL_NAME
+    return lora_path if lora_path.exists() else None
+
+
+def validate_image_file(filename: str) -> bool:
+    """
+    Validate if a file has an allowed image extension
+    
+    Args:
+        filename: Name of the file to validate
+        
+    Returns:
+        True if valid, False otherwise
+    """
+    return Path(filename).suffix.lower() in settings.ALLOWED_EXTENSIONS
+
+
+if __name__ == "__main__":
+    # Test configuration
+    ensure_directories()
+    print(f"\n{settings.APP_NAME} v{settings.APP_VERSION}")
+    print(f"Base Model: {settings.BASE_MODEL}")
+    print(f"LoRA Path: {get_lora_path()}")

backend/services/diffusion_pipeline.py

@@ -0,0 +1,340 @@
+"""
+Diffusion Pipeline Service
+
+Orchestrates the complete image enhancement pipeline:
+1. Load and preprocess input image
+2. Run Stable Diffusion img2img with optional LoRA
+3. Apply upscaling
+4. Apply post-processing effects
+"""
+import logging
+from pathlib import Path
+from typing import Optional, Callable
+import torch
+from PIL import Image
+from diffusers import (
+    StableDiffusionImg2ImgPipeline,
+    AutoencoderKL,
+    DPMSolverMultistepScheduler
+)
+
+from config import settings, get_lora_path
+from services.upscaler import upscale_image
+from services.postprocess import postprocess_image
+
+
+logger = logging.getLogger(__name__)
+
+
+class DiffusionPipelineManager:
+    """
+    Manages the Stable Diffusion pipeline for image enhancement
+    
+    This class handles lazy loading of models and provides a clean
+    interface for image-to-image enhancement.
+    """
+    
+    def __init__(self):
+        self.pipe = None
+        self.lora_loaded = False
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Pipeline will use device: {self.device}")
+        logger.info(f"CUDA available: {torch.cuda.is_available()}")
+        if torch.cuda.is_available():
+            logger.info(f"CUDA device: {torch.cuda.get_device_name(0)}")
+    
+    def load_pipeline(self):
+        """
+        Load the Stable Diffusion img2img pipeline with optimizations
+        
+        This method initializes the Stable Diffusion pipeline with:
+        - Mixed precision for faster inference
+        - Memory-efficient attention
+        - Optional VAE for better image quality
+        """
+        if self.pipe is not None:
+            return
+        
+        logger.info("Loading Stable Diffusion 1.5 pipeline...")
+        
+        try:
+            # Determine dtype based on device
+            dtype = torch.float32  # Always use float32 for CPU compatibility
+            
+            logger.info(f"Loading model with dtype: {dtype}, device: {self.device}")
+            
+            # Initialize the pipeline (SD 1.5) - let diffusers handle all components
+            self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                settings.BASE_MODEL,
+                torch_dtype=dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
+                low_cpu_mem_usage=False  # Load everything properly
+            )
+            
+            # Optimize scheduler
+            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+                self.pipe.scheduler.config
+            )
+            
+            # Move to device
+            self.pipe = self.pipe.to(self.device)
+            
+            # Enable memory optimizations
+            if self.device == "cuda":
+                if settings.ENABLE_ATTENTION_SLICING:
+                    self.pipe.enable_attention_slicing()
+                    logger.info("Enabled attention slicing")
+                if settings.ENABLE_VAE_SLICING:
+                    self.pipe.enable_vae_slicing()
+                    logger.info("Enabled VAE slicing")
+            
+            logger.info("Pipeline loaded successfully")
+            logger.info(f"Model: {settings.BASE_MODEL}")
+            logger.info(f"Device: {self.device}")
+            
+        except Exception as e:
+            logger.error(f"Failed to load pipeline: {e}", exc_info=True)
+            raise
+    
+    def load_lora(self, lora_path: Path):
+        """
+        Load a custom LoRA adapter into the pipeline
+        
+        Args:
+            lora_path: Path to the LoRA weights file (.safetensors)
+        """
+        if self.pipe is None:
+            self.load_pipeline()
+        
+        try:
+            logger.info(f"Loading LoRA from: {lora_path}")
+            self.pipe.load_lora_weights(str(lora_path.parent), weight_name=lora_path.name)
+            self.lora_loaded = True
+            logger.info("LoRA loaded successfully")
+            
+        except Exception as e:
+            logger.warning(f"Failed to load LoRA (continuing without it): {e}")
+            self.lora_loaded = False
+    
+    def enhance(
+        self,
+        image: Image.Image,
+        prompt: str,
+        negative_prompt: str,
+        strength: float,
+        guidance_scale: float,
+        num_inference_steps: int = 30,
+        progress_callback: Optional[Callable[[int, str], None]] = None
+    ) -> Image.Image:
+        """
+        Run image-to-image enhancement
+        
+        Args:
+            image: Input PIL Image
+            prompt: Positive prompt
+            negative_prompt: Negative prompt
+            strength: Denoising strength (0-1)
+            guidance_scale: CFG scale
+            num_inference_steps: Number of diffusion steps
+            progress_callback: Optional callback for progress updates
+            
+        Returns:
+            Enhanced PIL Image
+        """
+        if self.pipe is None:
+            if progress_callback:
+                progress_callback(30, "Loading AI model...")
+            self.load_pipeline()
+        
+        if progress_callback:
+            progress_callback(40, "Preprocessing image...")
+        
+        logger.info(f"Running enhancement - Strength: {strength}, Guidance: {guidance_scale}")
+        
+        # Ensure image is in RGB mode
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        
+        # Resize if too large
+        max_size = settings.MAX_IMAGE_SIZE
+        if max(image.size) > max_size:
+            ratio = max_size / max(image.size)
+            new_size = tuple(int(dim * ratio) for dim in image.size)
+            image = image.resize(new_size, Image.LANCZOS)
+            logger.info(f"Resized input to: {new_size}")
+        
+        if progress_callback:
+            progress_callback(50, f"Starting generation ({num_inference_steps} steps)...")
+        
+        # Callback to track diffusion progress
+        def step_callback(pipe, step_index, timestep, callback_kwargs):
+            if progress_callback:
+                # Calculate progress: 50-85% range mapped to steps
+                progress = 50 + int((step_index / num_inference_steps) * 35)
+                progress_callback(progress, f"Step {step_index+1}/{num_inference_steps}")
+            return callback_kwargs
+        
+        # Run inference
+        with torch.inference_mode():
+            result = self.pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=image,
+                strength=strength,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                callback_on_step_end=step_callback if progress_callback else None
+            ).images[0]
+        
+        if progress_callback:
+            progress_callback(85, "Finalizing...")
+        
+        return result
+    
+    def unload(self):
+        """
+        Unload the pipeline to free memory
+        """
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+            self.lora_loaded = False
+            torch.cuda.empty_cache()
+            logger.info("Pipeline unloaded")
+
+
+# Global pipeline instance (singleton pattern)
+_pipeline_manager = None
+
+
+def get_pipeline_manager() -> DiffusionPipelineManager:
+    """
+    Get or create the global pipeline manager instance
+    
+    Returns:
+        DiffusionPipelineManager instance
+    """
+    global _pipeline_manager
+    if _pipeline_manager is None:
+        _pipeline_manager = DiffusionPipelineManager()
+    return _pipeline_manager
+
+
+def preprocess_image(image_path: Path) -> Image.Image:
+    """
+    Load and preprocess an input image
+    
+    This function:
+    - Loads the image
+    - Converts to RGB
+    - Normalizes exposure/gamma if needed
+    
+    Args:
+        image_path: Path to input image
+        
+    Returns:
+        Preprocessed PIL Image
+    """
+    logger.info(f"Loading image from: {image_path}")
+    
+    image = Image.open(image_path)
+    
+    # Convert to RGB if necessary
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    # Optional: Add gamma/exposure normalization here
+    # This would analyze the histogram and adjust brightness
+    
+    return image
+
+
+def enhance_image(
+    input_image_path: Path,
+    output_path: Path,
+    strength: float = settings.IMG2IMG_STRENGTH,
+    guidance_scale: float = settings.GUIDANCE_SCALE,
+    custom_prompt: Optional[str] = None,
+    use_upscaler: bool = True,
+    use_postprocess: bool = True,
+    progress_callback: Optional[Callable[[int, str], None]] = None
+) -> Path:
+    """
+    Complete image enhancement pipeline
+    
+    This is the main entry point for enhancing images. It orchestrates:
+    1. Image preprocessing
+    2. Stable Diffusion img2img with optional LoRA
+    3. Upscaling (optional)
+    4. Post-processing (optional)
+    
+    Args:
+        input_image_path: Path to input image
+        output_path: Path to save enhanced image
+        strength: Image-to-image strength
+        guidance_scale: CFG scale
+        custom_prompt: Optional custom prompt
+        use_upscaler: Whether to apply upscaling
+        use_postprocess: Whether to apply post-processing
+        progress_callback: Optional callback for progress updates
+        
+    Returns:
+        Path to the saved enhanced image
+    """
+    try:
+        # Step 1: Load and preprocess
+        if progress_callback:
+            progress_callback(25, "Loading and preprocessing image...")
+        logger.info("Step 1/4: Preprocessing image")
+        image = preprocess_image(input_image_path)
+        
+        # Step 2: Get pipeline and load LoRA if available
+        logger.info("Step 2/4: Running diffusion model")
+        pipeline = get_pipeline_manager()
+        
+        lora_path = get_lora_path()
+        if lora_path and not pipeline.lora_loaded:
+            pipeline.load_lora(lora_path)
+        
+        # Determine prompt
+        prompt = custom_prompt if custom_prompt else settings.DEFAULT_PROMPT
+        
+        # Run enhancement
+        enhanced = pipeline.enhance(
+            image=image,
+            prompt=prompt,
+            negative_prompt=settings.NEGATIVE_PROMPT,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=settings.NUM_INFERENCE_STEPS,
+            progress_callback=progress_callback
+        )
+        
+        # Step 3: Upscaling
+        if use_upscaler:
+            if progress_callback:
+                progress_callback(92, "Upscaling image...")
+            logger.info("Step 3/4: Upscaling image")
+            enhanced = upscale_image(enhanced, scale=settings.UPSCALE_FACTOR)
+        else:
+            logger.info("Step 3/4: Skipping upscaling")
+        
+        # Step 4: Post-processing
+        if use_postprocess:
+            if progress_callback:
+                progress_callback(94, "Applying final touches...")
+            logger.info("Step 4/4: Applying post-processing")
+            enhanced = postprocess_image(enhanced)
+        else:
+            logger.info("Step 4/4: Skipping post-processing")
+        
+        # Save result
+        enhanced.save(output_path, quality=95, optimize=True)
+        logger.info(f"Enhanced image saved to: {output_path}")
+        
+        return output_path
+        
+    except Exception as e:
+        logger.error(f"Enhancement pipeline failed: {e}", exc_info=True)
+        raise

backend/services/postprocess.py

@@ -0,0 +1,294 @@
+"""
+Post-Processing Service
+
+Applies subtle photographic enhancements to the generated images:
+- Local contrast enhancement (CLAHE)
+- Film grain simulation
+- Subtle vignette effect
+- Color grading adjustments
+"""
+import logging
+import numpy as np
+from PIL import Image, ImageEnhance, ImageFilter
+import cv2
+
+
+logger = logging.getLogger(__name__)
+
+
+def apply_clahe(image: Image.Image, clip_limit: float = 2.0) -> Image.Image:
+    """
+    Apply Contrast Limited Adaptive Histogram Equalization (CLAHE)
+    
+    This enhances local contrast without over-amplifying noise.
+    
+    Args:
+        image: Input PIL Image
+        clip_limit: Threshold for contrast limiting (higher = more contrast)
+        
+    Returns:
+        Enhanced PIL Image
+    """
+    # Convert to numpy array
+    img_np = np.array(image)
+    
+    # Convert to LAB color space
+    lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
+    
+    # Split channels
+    l, a, b = cv2.split(lab)
+    
+    # Apply CLAHE to L channel
+    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
+    l_clahe = clahe.apply(l)
+    
+    # Merge channels
+    lab_clahe = cv2.merge([l_clahe, a, b])
+    
+    # Convert back to RGB
+    rgb = cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB)
+    
+    return Image.fromarray(rgb)
+
+
+def add_film_grain(
+    image: Image.Image,
+    intensity: float = 0.02,
+    grain_size: float = 1.0
+) -> Image.Image:
+    """
+    Add subtle film grain for a more organic look
+    
+    Args:
+        image: Input PIL Image
+        intensity: Strength of the grain effect (0.01-0.05 recommended)
+        grain_size: Size of grain particles
+        
+    Returns:
+        Image with film grain
+    """
+    img_np = np.array(image).astype(np.float32) / 255.0
+    
+    # Generate noise
+    noise = np.random.normal(0, intensity, img_np.shape)
+    
+    # Optional: blur noise for larger grain
+    if grain_size > 1.0:
+        noise = cv2.GaussianBlur(noise, (0, 0), grain_size)
+    
+    # Add noise to image
+    noisy = img_np + noise
+    noisy = np.clip(noisy, 0, 1)
+    
+    # Convert back to uint8
+    result = (noisy * 255).astype(np.uint8)
+    
+    return Image.fromarray(result)
+
+
+def apply_vignette(
+    image: Image.Image,
+    strength: float = 0.3,
+    radius: float = 0.8
+) -> Image.Image:
+    """
+    Apply a subtle vignette effect
+    
+    Darkens the corners and edges of the image to draw focus to the center.
+    
+    Args:
+        image: Input PIL Image
+        strength: Vignette intensity (0-1)
+        radius: Radius of the unaffected center area (0-1)
+        
+    Returns:
+        Image with vignette
+    """
+    width, height = image.size
+    img_np = np.array(image).astype(np.float32)
+    
+    # Create coordinate grids
+    x = np.linspace(-1, 1, width)
+    y = np.linspace(-1, 1, height)
+    X, Y = np.meshgrid(x, y)
+    
+    # Calculate distance from center
+    distance = np.sqrt(X**2 + Y**2)
+    
+    # Create vignette mask
+    vignette = 1 - np.clip((distance - radius) / (1 - radius), 0, 1) * strength
+    vignette = vignette[:, :, np.newaxis]  # Add channel dimension
+    
+    # Apply vignette
+    result = img_np * vignette
+    result = np.clip(result, 0, 255).astype(np.uint8)
+    
+    return Image.fromarray(result)
+
+
+def enhance_colors(
+    image: Image.Image,
+    saturation: float = 1.1,
+    contrast: float = 1.05,
+    brightness: float = 1.0
+) -> Image.Image:
+    """
+    Apply subtle color grading adjustments
+    
+    Args:
+        image: Input PIL Image
+        saturation: Saturation multiplier (1.0 = no change)
+        contrast: Contrast multiplier (1.0 = no change)
+        brightness: Brightness multiplier (1.0 = no change)
+        
+    Returns:
+        Color-graded image
+    """
+    # Adjust saturation
+    if saturation != 1.0:
+        enhancer = ImageEnhance.Color(image)
+        image = enhancer.enhance(saturation)
+    
+    # Adjust contrast
+    if contrast != 1.0:
+        enhancer = ImageEnhance.Contrast(image)
+        image = enhancer.enhance(contrast)
+    
+    # Adjust brightness
+    if brightness != 1.0:
+        enhancer = ImageEnhance.Brightness(image)
+        image = enhancer.enhance(brightness)
+    
+    return image
+
+
+def sharpen_image(image: Image.Image, strength: float = 1.0) -> Image.Image:
+    """
+    Apply subtle sharpening
+    
+    Args:
+        image: Input PIL Image
+        strength: Sharpening strength (0-2 recommended)
+        
+    Returns:
+        Sharpened image
+    """
+    if strength <= 0:
+        return image
+    
+    # Use UnsharpMask for better control
+    from PIL import ImageFilter
+    
+    # Blend between original and sharpened
+    sharpened = image.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
+    
+    if strength < 1.0:
+        # Blend with original
+        return Image.blend(image, sharpened, strength)
+    else:
+        return sharpened
+
+
+def postprocess_image(
+    image: Image.Image,
+    apply_contrast: bool = True,
+    apply_grain: bool = True,
+    apply_vignette_effect: bool = True,
+    apply_color_grading: bool = True,
+    apply_sharpening: bool = True
+) -> Image.Image:
+    """
+    Apply complete post-processing pipeline
+    
+    This function orchestrates all post-processing effects in the optimal order:
+    1. Local contrast enhancement (CLAHE)
+    2. Color grading
+    3. Sharpening
+    4. Film grain
+    5. Vignette
+    
+    Args:
+        image: Input PIL Image
+        apply_contrast: Enable local contrast enhancement
+        apply_grain: Enable film grain
+        apply_vignette_effect: Enable vignette
+        apply_color_grading: Enable color adjustments
+        apply_sharpening: Enable sharpening
+        
+    Returns:
+        Post-processed PIL Image
+    """
+    logger.info("Starting post-processing")
+    
+    try:
+        # Step 1: Local contrast
+        if apply_contrast:
+            logger.debug("Applying CLAHE")
+            image = apply_clahe(image, clip_limit=2.0)
+        
+        # Step 2: Color grading
+        if apply_color_grading:
+            logger.debug("Applying color grading")
+            image = enhance_colors(
+                image,
+                saturation=1.08,  # Slightly more saturated
+                contrast=1.03,     # Slightly more contrast
+                brightness=1.0     # No brightness change
+            )
+        
+        # Step 3: Sharpening
+        if apply_sharpening:
+            logger.debug("Applying sharpening")
+            image = sharpen_image(image, strength=0.6)
+        
+        # Step 4: Film grain
+        if apply_grain:
+            logger.debug("Adding film grain")
+            image = add_film_grain(image, intensity=0.015, grain_size=1.2)
+        
+        # Step 5: Vignette
+        if apply_vignette_effect:
+            logger.debug("Applying vignette")
+            image = apply_vignette(image, strength=0.2, radius=0.85)
+        
+        logger.info("Post-processing completed")
+        return image
+        
+    except Exception as e:
+        logger.error(f"Error in post-processing: {e}", exc_info=True)
+        logger.warning("Returning original image")
+        return image
+
+
+def create_comparison(
+    original: Image.Image,
+    processed: Image.Image,
+    padding: int = 10
+) -> Image.Image:
+    """
+    Create a side-by-side comparison image
+    
+    Useful for visualizing before/after results.
+    
+    Args:
+        original: Original image
+        processed: Processed image
+        padding: Space between images in pixels
+        
+    Returns:
+        Combined comparison image
+    """
+    # Ensure both images are the same size
+    if original.size != processed.size:
+        processed = processed.resize(original.size, Image.LANCZOS)
+    
+    width, height = original.size
+    
+    # Create new image with space for both
+    comparison = Image.new('RGB', (width * 2 + padding, height), color='white')
+    
+    # Paste images
+    comparison.paste(original, (0, 0))
+    comparison.paste(processed, (width + padding, 0))
+    
+    return comparison

backend/services/progress_tracker.py

@@ -0,0 +1,146 @@
+"""
+Progress Tracking System
+
+Manages progress updates for long-running enhancement tasks
+Supports Server-Sent Events (SSE) for real-time frontend updates
+"""
+import asyncio
+import logging
+from typing import Dict, Optional
+from datetime import datetime
+
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressTracker:
+    """
+    Tracks progress for enhancement operations
+    Allows multiple clients to subscribe to progress updates
+    """
+    
+    def __init__(self):
+        self.tasks: Dict[str, Dict] = {}
+        self.subscribers: Dict[str, list] = {}
+    
+    def create_task(self, task_id: str, total_steps: int = 100):
+        """Initialize a new task for progress tracking"""
+        self.tasks[task_id] = {
+            "id": task_id,
+            "status": "initializing",
+            "progress": 0,
+            "total_steps": total_steps,
+            "current_step": 0,
+            "message": "Starting enhancement...",
+            "started_at": datetime.now().isoformat(),
+            "error": None
+        }
+        self.subscribers[task_id] = []
+        logger.info(f"Created progress tracker for task {task_id}")
+    
+    def update_progress(
+        self, 
+        task_id: str, 
+        current_step: int, 
+        message: str,
+        status: str = "processing"
+    ):
+        """Update progress for a task (synchronous)"""
+        if task_id not in self.tasks:
+            logger.warning(f"Task {task_id} not found for progress update")
+            return
+        
+        task = self.tasks[task_id]
+        task["current_step"] = current_step
+        task["progress"] = int((current_step / task["total_steps"]) * 100)
+        task["message"] = message
+        task["status"] = status
+        
+        logger.info(f"Task {task_id}: {task['progress']}% - {message}")
+    
+    def complete_task(self, task_id: str, result_url: Optional[str] = None):
+        """Mark a task as completed"""
+        if task_id not in self.tasks:
+            return
+        
+        self.tasks[task_id].update({
+            "status": "completed",
+            "progress": 100,
+            "message": "Enhancement completed!",
+            "result_url": result_url,
+            "completed_at": datetime.now().isoformat()
+        })
+        
+        # Notify subscribers if event loop is running
+        try:
+            asyncio.create_task(self._notify_subscribers(task_id))
+        except RuntimeError:
+            # No event loop running (called from thread)
+            pass
+        
+        logger.info(f"Task {task_id} completed")
+    
+    def fail_task(self, task_id: str, error: str):
+        """Mark a task as failed"""
+        if task_id not in self.tasks:
+            return
+        
+        self.tasks[task_id].update({
+            "status": "failed",
+            "message": f"Error: {error}",
+            "error": error,
+            "failed_at": datetime.now().isoformat()
+        })
+        
+        # Notify subscribers if event loop is running
+        try:
+            asyncio.create_task(self._notify_subscribers(task_id))
+        except RuntimeError:
+            # No event loop running (called from thread)
+            pass
+        
+        logger.error(f"Task {task_id} failed: {error}")
+    
+    def get_task_status(self, task_id: str) -> Optional[Dict]:
+        """Get current status of a task"""
+        return self.tasks.get(task_id)
+    
+    def subscribe(self, task_id: str, queue: asyncio.Queue):
+        """Subscribe to progress updates for a task"""
+        if task_id not in self.subscribers:
+            self.subscribers[task_id] = []
+        self.subscribers[task_id].append(queue)
+        logger.debug(f"New subscriber for task {task_id}")
+    
+    def unsubscribe(self, task_id: str, queue: asyncio.Queue):
+        """Unsubscribe from progress updates"""
+        if task_id in self.subscribers and queue in self.subscribers[task_id]:
+            self.subscribers[task_id].remove(queue)
+    
+    async def _notify_subscribers(self, task_id: str):
+        """Send current task status to all subscribers"""
+        if task_id not in self.subscribers:
+            return
+        
+        task_data = self.tasks.get(task_id)
+        if not task_data:
+            return
+        
+        # Send to all subscribers
+        for queue in self.subscribers[task_id]:
+            try:
+                await queue.put(task_data.copy())
+            except Exception as e:
+                logger.error(f"Error notifying subscriber: {e}")
+    
+    def cleanup_task(self, task_id: str):
+        """Remove task data after completion"""
+        if task_id in self.tasks:
+            del self.tasks[task_id]
+        if task_id in self.subscribers:
+            del self.subscribers[task_id]
+        logger.debug(f"Cleaned up task {task_id}")
+
+
+# Global progress tracker instance
+progress_tracker = ProgressTracker()

backend/services/training_engine.py

@@ -0,0 +1,369 @@
+"""
+Training Engine for LoRA Models
+
+This module handles the complete training pipeline for custom LoRA adapters:
+- Dataset preparation and loading
+- Model initialization
+- Training loop with checkpointing
+- Model saving and validation
+"""
+import logging
+import os
+from pathlib import Path
+from typing import Optional
+import torch
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+from diffusers import (
+    StableDiffusionXLPipeline,
+    AutoencoderKL,
+    UNet2DConditionModel
+)
+from transformers import CLIPTextModel, CLIPTokenizer
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from peft import LoraConfig, get_peft_model
+import torch.nn.functional as F
+
+from config import settings
+
+
+logger = logging.getLogger(__name__)
+
+
+class ArchitectureDataset(Dataset):
+    """
+    Custom dataset for architectural image pairs
+    
+    Loads input-target pairs for LoRA training.
+    """
+    
+    def __init__(self, input_dir: Path, target_dir: Path, size: int = 1024):
+        """
+        Initialize dataset
+        
+        Args:
+            input_dir: Directory containing input images
+            target_dir: Directory containing target images
+            size: Target image size
+        """
+        self.input_dir = input_dir
+        self.target_dir = target_dir
+        self.size = size
+        
+        # Load all pairs
+        self.pairs = self._load_pairs()
+        
+        logger.info(f"Loaded {len(self.pairs)} training pairs")
+    
+    def _load_pairs(self):
+        """
+        Match input and target images by their pair ID
+        
+        Returns:
+            List of (input_path, target_path) tuples
+        """
+        pairs = []
+        
+        # Get all input files
+        input_files = sorted(self.input_dir.glob("*_input.*"))
+        
+        for input_file in input_files:
+            # Extract pair ID
+            pair_id = input_file.stem.replace("_input", "")
+            
+            # Find matching target file
+            target_files = list(self.target_dir.glob(f"{pair_id}_target.*"))
+            
+            if target_files:
+                pairs.append((input_file, target_files[0]))
+            else:
+                logger.warning(f"No target found for input: {input_file}")
+        
+        return pairs
+    
+    def __len__(self):
+        return len(self.pairs)
+    
+    def __getitem__(self, idx):
+        """
+        Get a training pair
+        
+        Returns:
+            Dictionary with 'input' and 'target' tensors
+        """
+        input_path, target_path = self.pairs[idx]
+        
+        # Load images
+        input_img = Image.open(input_path).convert("RGB")
+        target_img = Image.open(target_path).convert("RGB")
+        
+        # Resize to target size
+        input_img = input_img.resize((self.size, self.size), Image.LANCZOS)
+        target_img = target_img.resize((self.size, self.size), Image.LANCZOS)
+        
+        # Convert to tensors and normalize to [-1, 1]
+        input_tensor = torch.from_numpy(
+            (torch.tensor(input_img).float() / 127.5 - 1.0).numpy()
+        ).permute(2, 0, 1)
+        
+        target_tensor = torch.from_numpy(
+            (torch.tensor(target_img).float() / 127.5 - 1.0).numpy()
+        ).permute(2, 0, 1)
+        
+        return {
+            "input": input_tensor,
+            "target": target_tensor,
+            "input_path": str(input_path),
+            "target_path": str(target_path)
+        }
+
+
+def prepare_lora_config(rank: int = 8) -> LoraConfig:
+    """
+    Create LoRA configuration
+    
+    Args:
+        rank: LoRA rank (dimensionality of adaptation matrices)
+        
+    Returns:
+        LoraConfig object
+    """
+    return LoraConfig(
+        r=rank,
+        lora_alpha=rank,  # Often set equal to rank
+        target_modules=[
+            "to_q", "to_k", "to_v", "to_out.0",  # Attention layers
+            "proj_in", "proj_out",                 # Projections
+            "ff.net.0.proj", "ff.net.2"           # Feed-forward layers
+        ],
+        lora_dropout=0.1,
+        bias="none",
+        task_type="CAUSAL_LM"
+    )
+
+
+def train_lora_model(
+    train_steps: int = settings.TRAIN_STEPS,
+    learning_rate: float = settings.LEARNING_RATE,
+    lora_rank: int = settings.LORA_RANK,
+    batch_size: int = settings.BATCH_SIZE,
+    lock_file: Optional[Path] = None
+):
+    """
+    Main training function for LoRA model
+    
+    This function orchestrates the complete training process:
+    1. Initialize accelerator and models
+    2. Prepare dataset and dataloader
+    3. Training loop with gradient accumulation
+    4. Save final model
+    
+    Args:
+        train_steps: Total number of training steps
+        learning_rate: Learning rate for optimizer
+        lora_rank: Rank for LoRA adaptation
+        batch_size: Batch size for training
+        lock_file: Optional lock file to remove after training
+    """
+    try:
+        logger.info("=" * 60)
+        logger.info("Starting LoRA Training")
+        logger.info("=" * 60)
+        logger.info(f"Configuration:")
+        logger.info(f"  - Train Steps: {train_steps}")
+        logger.info(f"  - Learning Rate: {learning_rate}")
+        logger.info(f"  - LoRA Rank: {lora_rank}")
+        logger.info(f"  - Batch Size: {batch_size}")
+        logger.info(f"  - Base Model: {settings.BASE_MODEL}")
+        
+        # Initialize accelerator
+        accelerator = Accelerator(
+            mixed_precision=settings.MIXED_PRECISION,
+            gradient_accumulation_steps=settings.GRADIENT_ACCUMULATION_STEPS,
+            log_with="tensorboard",
+            project_dir=str(settings.LOGS_DIR)
+        )
+        
+        # Set random seed for reproducibility
+        set_seed(42)
+        
+        # Load dataset
+        logger.info("Loading training dataset...")
+        dataset = ArchitectureDataset(
+            input_dir=settings.INPUT_DIR,
+            target_dir=settings.TARGET_DIR,
+            size=settings.TARGET_RESOLUTION
+        )
+        
+        if len(dataset) == 0:
+            raise ValueError("No training data found!")
+        
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,  # Set to 0 for Windows compatibility
+            pin_memory=True
+        )
+        
+        # Load base model components
+        logger.info("Loading Stable Diffusion XL components...")
+        
+        # Load UNet (the main model we'll adapt)
+        unet = UNet2DConditionModel.from_pretrained(
+            settings.BASE_MODEL,
+            subfolder="unet",
+            torch_dtype=torch.float16 if accelerator.device.type == "cuda" else torch.float32
+        )
+        
+        # Load VAE for encoding images
+        vae = AutoencoderKL.from_pretrained(
+            settings.BASE_MODEL,
+            subfolder="vae",
+            torch_dtype=torch.float16 if accelerator.device.type == "cuda" else torch.float32
+        )
+        vae.requires_grad_(False)  # Freeze VAE
+        
+        # Load text encoder
+        text_encoder = CLIPTextModel.from_pretrained(
+            settings.BASE_MODEL,
+            subfolder="text_encoder"
+        )
+        text_encoder.requires_grad_(False)  # Freeze text encoder
+        
+        tokenizer = CLIPTokenizer.from_pretrained(
+            settings.BASE_MODEL,
+            subfolder="tokenizer"
+        )
+        
+        # Apply LoRA to UNet
+        logger.info(f"Applying LoRA with rank {lora_rank}...")
+        lora_config = prepare_lora_config(rank=lora_rank)
+        unet = get_peft_model(unet, lora_config)
+        unet.print_trainable_parameters()
+        
+        # Setup optimizer
+        optimizer = torch.optim.AdamW(
+            unet.parameters(),
+            lr=learning_rate,
+            betas=(0.9, 0.999),
+            weight_decay=1e-2,
+            eps=1e-8
+        )
+        
+        # Prepare with accelerator
+        unet, optimizer, dataloader = accelerator.prepare(
+            unet, optimizer, dataloader
+        )
+        vae = vae.to(accelerator.device)
+        text_encoder = text_encoder.to(accelerator.device)
+        
+        # Training loop
+        logger.info("Starting training loop...")
+        global_step = 0
+        progress_bar = range(train_steps)
+        
+        # Encode prompt for all samples
+        prompt = settings.DEFAULT_PROMPT
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt"
+        )
+        text_embeddings = text_encoder(text_inputs.input_ids.to(accelerator.device))[0]
+        
+        unet.train()
+        
+        for epoch in range(100):  # Large number, will break when steps reached
+            for batch in dataloader:
+                with accelerator.accumulate(unet):
+                    # Encode images to latent space
+                    with torch.no_grad():
+                        latents_input = vae.encode(batch["input"].to(accelerator.device)).latent_dist.sample()
+                        latents_target = vae.encode(batch["target"].to(accelerator.device)).latent_dist.sample()
+                        
+                        # Scale latents
+                        latents_input = latents_input * vae.config.scaling_factor
+                        latents_target = latents_target * vae.config.scaling_factor
+                    
+                    # Sample random timestep
+                    timesteps = torch.randint(
+                        0, 1000, (latents_input.shape[0],),
+                        device=accelerator.device
+                    ).long()
+                    
+                    # Add noise to target latents
+                    noise = torch.randn_like(latents_target)
+                    noisy_latents = latents_target  # Simplified for img2img
+                    
+                    # Predict noise
+                    model_pred = unet(
+                        noisy_latents,
+                        timesteps,
+                        text_embeddings.repeat(latents_input.shape[0], 1, 1)
+                    ).sample
+                    
+                    # Calculate loss (MSE between prediction and target)
+                    loss = F.mse_loss(model_pred.float(), latents_target.float(), reduction="mean")
+                    
+                    # Backpropagation
+                    accelerator.backward(loss)
+                    
+                    if accelerator.sync_gradients:
+                        accelerator.clip_grad_norm_(unet.parameters(), settings.MAX_GRAD_NORM)
+                    
+                    optimizer.step()
+                    optimizer.zero_grad()
+                
+                if accelerator.sync_gradients:
+                    global_step += 1
+                    
+                    if global_step % 50 == 0:
+                        logger.info(f"Step {global_step}/{train_steps} - Loss: {loss.item():.4f}")
+                    
+                    # Save checkpoint
+                    if global_step % settings.SAVE_STEPS == 0:
+                        logger.info(f"Saving checkpoint at step {global_step}")
+                        save_path = settings.LORA_DIR / f"checkpoint_{global_step}"
+                        accelerator.unwrap_model(unet).save_pretrained(save_path)
+                    
+                    # Check if training complete
+                    if global_step >= train_steps:
+                        break
+            
+            if global_step >= train_steps:
+                break
+        
+        # Save final model
+        logger.info("Training completed! Saving final model...")
+        final_path = settings.LORA_DIR / settings.LORA_MODEL_NAME.replace(".safetensors", "")
+        accelerator.unwrap_model(unet).save_pretrained(final_path)
+        
+        # Also save as safetensors
+        from safetensors.torch import save_file
+        state_dict = accelerator.unwrap_model(unet).state_dict()
+        safetensors_path = settings.LORA_DIR / settings.LORA_MODEL_NAME
+        save_file(state_dict, safetensors_path)
+        
+        logger.info(f"Model saved to: {safetensors_path}")
+        logger.info("=" * 60)
+        logger.info("Training Complete!")
+        logger.info("=" * 60)
+        
+    except Exception as e:
+        logger.error(f"Training failed: {e}", exc_info=True)
+        raise
+    
+    finally:
+        # Remove lock file
+        if lock_file and lock_file.exists():
+            lock_file.unlink()
+            logger.info("Training lock removed")
+
+
+if __name__ == "__main__":
+    # Test training
+    train_lora_model(train_steps=100)

backend/services/upscaler.py

@@ -0,0 +1,197 @@
+"""
+Upscaler Service
+
+Provides image upscaling functionality using Real-ESRGAN or similar models.
+This enhances the resolution of generated images for higher quality output.
+"""
+import logging
+from typing import Optional
+import numpy as np
+from PIL import Image
+import cv2
+
+
+logger = logging.getLogger(__name__)
+
+
+# Global upscaler instance
+_upscaler = None
+
+
+def get_upscaler():
+    """
+    Get or create the upscaler model instance
+    
+    This uses Real-ESRGAN for high-quality upscaling.
+    Falls back to simple interpolation if Real-ESRGAN is not available.
+    
+    Returns:
+        Upscaler instance or None
+    """
+    global _upscaler
+    
+    if _upscaler is not None:
+        return _upscaler
+    
+    try:
+        from realesrgan import RealESRGANer
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        
+        logger.info("Loading Real-ESRGAN model...")
+        
+        # Initialize the model
+        # Using RealESRGAN_x4plus for 4x upscaling
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=4
+        )
+        
+        # You can download the model weights from:
+        # https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth
+        model_path = "models/base/RealESRGAN_x4plus.pth"
+        
+        _upscaler = RealESRGANer(
+            scale=4,
+            model_path=model_path,
+            model=model,
+            tile=400,  # Tile size for processing large images
+            tile_pad=10,
+            pre_pad=0,
+            half=True  # Use FP16 for faster inference
+        )
+        
+        logger.info("Real-ESRGAN loaded successfully")
+        return _upscaler
+        
+    except ImportError:
+        logger.warning("Real-ESRGAN not available. Will use fallback upscaling.")
+        return None
+    except Exception as e:
+        logger.warning(f"Failed to load Real-ESRGAN: {e}. Using fallback.")
+        return None
+
+
+def upscale_with_realesrgan(image: Image.Image, scale: int = 2) -> Image.Image:
+    """
+    Upscale image using Real-ESRGAN
+    
+    Args:
+        image: Input PIL Image
+        scale: Upscaling factor
+        
+    Returns:
+        Upscaled PIL Image
+    """
+    upscaler = get_upscaler()
+    
+    if upscaler is None:
+        raise RuntimeError("Real-ESRGAN not available")
+    
+    # Convert PIL to numpy array (BGR for OpenCV)
+    img_np = np.array(image)
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    
+    # Upscale
+    output, _ = upscaler.enhance(img_bgr, outscale=scale)
+    
+    # Convert back to PIL (RGB)
+    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+    result = Image.fromarray(output_rgb)
+    
+    return result
+
+
+def upscale_with_lanczos(image: Image.Image, scale: int = 2) -> Image.Image:
+    """
+    Fallback upscaling using Lanczos interpolation
+    
+    This is a simple but effective upscaling method when Real-ESRGAN
+    is not available.
+    
+    Args:
+        image: Input PIL Image
+        scale: Upscaling factor
+        
+    Returns:
+        Upscaled PIL Image
+    """
+    width, height = image.size
+    new_size = (width * scale, height * scale)
+    
+    logger.info(f"Upscaling with Lanczos: {image.size} -> {new_size}")
+    
+    return image.resize(new_size, Image.LANCZOS)
+
+
+def upscale_image(image: Image.Image, scale: int = 2) -> Image.Image:
+    """
+    Upscale an image using the best available method
+    
+    Tries Real-ESRGAN first, falls back to Lanczos if unavailable.
+    
+    Args:
+        image: Input PIL Image
+        scale: Upscaling factor (2 or 4 recommended)
+        
+    Returns:
+        Upscaled PIL Image
+    """
+    try:
+        # Try Real-ESRGAN
+        logger.info(f"Upscaling image by {scale}x")
+        result = upscale_with_realesrgan(image, scale=scale)
+        logger.info("Upscaling completed with Real-ESRGAN")
+        return result
+        
+    except (RuntimeError, Exception) as e:
+        # Fallback to Lanczos
+        logger.info(f"Using fallback upscaling method: {e}")
+        return upscale_with_lanczos(image, scale=scale)
+
+
+def adaptive_upscale(
+    image: Image.Image,
+    target_size: Optional[int] = None,
+    max_scale: int = 4
+) -> Image.Image:
+    """
+    Upscale image adaptively to reach a target size
+    
+    This function calculates the appropriate scale factor to reach
+    the target size without exceeding max_scale.
+    
+    Args:
+        image: Input PIL Image
+        target_size: Target resolution for the longest edge
+        max_scale: Maximum upscaling factor
+        
+    Returns:
+        Upscaled PIL Image
+    """
+    if target_size is None:
+        target_size = 2048  # Default target
+    
+    current_size = max(image.size)
+    
+    if current_size >= target_size:
+        logger.info("Image already at target size or larger")
+        return image
+    
+    # Calculate required scale
+    required_scale = target_size / current_size
+    
+    # Clamp to max_scale and round to nearest power of 2
+    if required_scale <= 1:
+        scale = 1
+    elif required_scale <= 2:
+        scale = 2
+    else:
+        scale = min(4, max_scale)
+    
+    logger.info(f"Adaptive upscaling: {scale}x (current: {current_size}px, target: {target_size}px)")
+    
+    return upscale_image(image, scale=scale)

requirements.txt

@@ -0,0 +1,27 @@
+# Requirements for Hugging Face Spaces deployment
+
+# Core ML libraries
+torch>=2.0.0
+torchvision>=0.15.0
+diffusers>=0.25.0
+transformers>=4.36.0
+accelerate>=0.25.0
+
+# Image processing
+Pillow>=10.0.0
+opencv-python-headless>=4.8.0
+numpy>=1.24.0
+
+# Upscaling
+basicsr>=1.4.2
+
+# Gradio interface
+gradio>=4.44.0
+
+# Configuration
+pydantic>=2.5.0
+pydantic-settings>=2.1.0
+
+# Utilities
+tqdm>=4.66.0
+safetensors>=0.4.0