Spaces:
Runtime error
Runtime error
| """ | |
| Architecture AI Enhancer - Simple Gradio Interface for HF Spaces | |
| """ | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| import logging | |
| from pathlib import Path | |
| import sys | |
| import os | |
| # Add backend to path | |
| sys.path.insert(0, str(Path(__file__).parent)) | |
| from backend.services.diffusion_pipeline import get_pipeline_manager | |
| from backend.services.upscaler import upscale_image | |
| from backend.services.postprocess import postprocess_image | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Initialize pipeline manager | |
| pipeline_manager = None | |
| def initialize_models(): | |
| """Initialize models on first use""" | |
| global pipeline_manager | |
| if pipeline_manager is None: | |
| logger.info("Loading Stable Diffusion pipeline...") | |
| pipeline_manager = get_pipeline_manager() | |
| # Check if there's a LoRA model to load | |
| from backend.config import get_lora_path | |
| lora_path = get_lora_path() | |
| if lora_path: | |
| logger.info(f"LoRA model found: {lora_path}") | |
| try: | |
| pipeline_manager.load_lora(lora_path) | |
| logger.info("✅ Custom LoRA loaded successfully!") | |
| except Exception as e: | |
| logger.warning(f"Failed to load LoRA: {e}") | |
| def enhance_image_simple( | |
| input_image: Image.Image, | |
| strength: float, | |
| guidance_scale: float, | |
| custom_prompt: str, | |
| use_upscaler: bool, | |
| use_postprocess: bool, | |
| progress=gr.Progress() | |
| ) -> Image.Image: | |
| """Enhance architectural image""" | |
| try: | |
| if input_image is None: | |
| raise ValueError("Please upload an image first") | |
| progress(0, desc="Starting enhancement...") | |
| # Initialize models | |
| try: | |
| progress(0.1, desc="Initializing AI models...") | |
| initialize_models() | |
| except Exception as e: | |
| logger.error(f"Model initialization failed: {e}", exc_info=True) | |
| raise ValueError(f"Failed to load AI models: {str(e)}") | |
| # Build prompt | |
| base_prompt = "professional architectural photography, highly detailed, 8k, photorealistic" | |
| if custom_prompt and custom_prompt.strip(): | |
| final_prompt = f"{base_prompt}, {custom_prompt.strip()}" | |
| else: | |
| final_prompt = base_prompt | |
| # Step 1: AI Enhancement | |
| progress(0.3, desc="Applying AI enhancement (this may take several minutes on CPU)...") | |
| try: | |
| # Create a callback to update progress during diffusion | |
| def update_progress(percent, description): | |
| # Map internal progress (30-85%) to our progress range | |
| mapped_progress = 0.3 + (percent / 100) * 0.4 # 30% to 70% | |
| progress(mapped_progress, desc=description) | |
| enhanced = pipeline_manager.enhance( | |
| image=input_image, | |
| prompt=final_prompt, | |
| negative_prompt="blurry, low quality, distorted, ugly, bad architecture", | |
| strength=strength, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=30, | |
| progress_callback=update_progress | |
| ) | |
| except Exception as e: | |
| logger.error(f"AI enhancement failed: {e}", exc_info=True) | |
| raise ValueError(f"AI enhancement failed: {str(e)}") | |
| # Step 2: Upscaling | |
| if use_upscaler: | |
| progress(0.75, desc="Upscaling image...") | |
| try: | |
| enhanced = upscale_image(enhanced, scale=2) | |
| except Exception as e: | |
| logger.warning(f"Upscaling failed, skipping: {e}") | |
| # Step 3: Post-processing | |
| if use_postprocess: | |
| progress(0.85, desc="Applying final touches...") | |
| try: | |
| enhanced = postprocess_image( | |
| enhanced, | |
| apply_contrast=True, | |
| apply_sharpening=True, | |
| apply_color_grading=True, | |
| apply_grain=False, | |
| apply_vignette_effect=False | |
| ) | |
| except Exception as e: | |
| logger.warning(f"Post-processing failed, skipping: {e}") | |
| progress(1.0, desc="Complete!") | |
| return enhanced | |
| except ValueError as e: | |
| # Re-raise ValueError with user-friendly message | |
| raise gr.Error(str(e)) | |
| except Exception as e: | |
| logger.error(f"Enhancement failed: {e}", exc_info=True) | |
| raise gr.Error(f"Enhancement failed: {str(e)}") | |
| # Create interface with tabs | |
| with gr.Blocks(title="🏛️ Architecture AI Enhancer") as demo: | |
| gr.Markdown(""" | |
| # 🏛️ Architecture AI Enhancer | |
| Transform your architectural renders with AI-powered enhancement using Stable Diffusion 1.5 | |
| """) | |
| # Check for custom LoRA | |
| from backend.config import get_lora_path | |
| lora_path = get_lora_path() | |
| if lora_path and lora_path.exists(): | |
| gr.Markdown(f""" | |
| <div style="background: linear-gradient(90deg, #4CAF50 0%, #45a049 100%); padding: 10px; border-radius: 8px; margin: 10px 0;"> | |
| <p style="color: white; margin: 0; font-weight: bold;"> | |
| ✅ Custom LoRA Active: {lora_path.name} | |
| </p> | |
| <p style="color: rgba(255,255,255,0.9); margin: 5px 0 0 0; font-size: 0.9em;"> | |
| All enhancements will use your custom trained style | |
| </p> | |
| </div> | |
| """) | |
| else: | |
| gr.Markdown(""" | |
| <div style="background: #f0f0f0; padding: 10px; border-radius: 8px; margin: 10px 0;"> | |
| <p style="color: #666; margin: 0;"> | |
| ℹ️ Using base Stable Diffusion 1.5 model. <a href="#" style="color: #2196F3;">Train a custom LoRA</a> for personalized style. | |
| </p> | |
| </div> | |
| """) | |
| with gr.Tabs(): | |
| # Tab 1: Enhancement | |
| with gr.Tab("✨ Enhance Image"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="📤 Input Image", type="pil") | |
| with gr.Accordion("⚙️ Advanced Settings", open=False): | |
| strength = gr.Slider( | |
| 0.1, 0.8, value=0.3, step=0.05, | |
| label="Denoising Strength", | |
| info="Lower = more faithful to input, Higher = more creative" | |
| ) | |
| guidance_scale = gr.Slider( | |
| 1.0, 15.0, value=5.5, step=0.5, | |
| label="Guidance Scale", | |
| info="How closely to follow the prompt" | |
| ) | |
| custom_prompt = gr.Textbox( | |
| label="Additional Prompt (Optional)", | |
| placeholder="e.g., modern minimalist, glass facade, sunset lighting...", | |
| lines=2, | |
| info="Add custom details to enhance specific aspects" | |
| ) | |
| use_upscaler = gr.Checkbox(label="Enable Upscaling (2x)", value=True) | |
| use_postprocess = gr.Checkbox(label="Enable Post-Processing", value=True) | |
| enhance_btn = gr.Button("✨ Enhance Image", variant="primary", size="lg") | |
| with gr.Column(): | |
| output_image = gr.Image(label="✅ Enhanced Result", type="pil") | |
| 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 | |
| """) | |
| # Connect the button | |
| enhance_btn.click( | |
| fn=enhance_image_simple, | |
| inputs=[input_image, strength, guidance_scale, custom_prompt, use_upscaler, use_postprocess], | |
| outputs=output_image | |
| ) | |
| # Tab 2: Training Guide | |
| with gr.Tab("🎓 Custom Training"): | |
| gr.Markdown(""" | |
| # 🎓 Train Your Own Custom LoRA Model | |
| Want to teach the AI your specific architectural style? You can train a custom LoRA model! | |
| ## ⚠️ Important Note | |
| **Training execution is not available on this HF Space** due to computational requirements. However, you can: | |
| 1. Upload and preview your training pairs below | |
| 2. Download the guide to train locally | |
| 3. Deploy your custom model | |
| """) | |
| gr.Markdown("---") | |
| # Training pair upload section | |
| gr.Markdown(""" | |
| ## 📸 Preview Training Pairs | |
| Upload example pairs to visualize what you'll train with: | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### Input Image (Before)") | |
| training_input = gr.Image(label="Input Render", type="pil") | |
| with gr.Column(): | |
| gr.Markdown("### Target Image (After)") | |
| training_target = gr.Image(label="Target Enhanced", type="pil") | |
| with gr.Row(): | |
| preview_btn = gr.Button("👁️ Preview Training Pair", variant="secondary") | |
| clear_btn = gr.Button("🗑️ Clear", variant="secondary") | |
| training_preview = gr.Gallery( | |
| label="Training Pair Preview", | |
| columns=2, | |
| rows=1, | |
| height="auto" | |
| ) | |
| def preview_training_pair(input_img, target_img): | |
| """Preview training pair side by side""" | |
| if input_img is None or target_img is None: | |
| return [] | |
| return [input_img, target_img] | |
| def clear_training(): | |
| """Clear training inputs""" | |
| return None, None, [] | |
| preview_btn.click( | |
| fn=preview_training_pair, | |
| inputs=[training_input, training_target], | |
| outputs=training_preview | |
| ) | |
| clear_btn.click( | |
| fn=clear_training, | |
| inputs=[], | |
| outputs=[training_input, training_target, training_preview] | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| ## 💡 Tips for Good Training Pairs | |
| - **Consistency**: Use similar compositions between input and target | |
| - **Quality**: High-resolution images (512x512 minimum) | |
| - **Variety**: Include different views and lighting conditions | |
| - **Alignment**: Input and target should show the same scene | |
| - **Target Quality**: Ensure targets represent your desired style accurately | |
| --- | |
| ## 🎯 Using a Custom LoRA Model | |
| Once you've trained a LoRA model locally, here's how to use it: | |
| ### Option 1: Deploy to This Space (Recommended) | |
| 1. **Upload your LoRA to Hugging Face Hub**: | |
| ```bash | |
| # Install huggingface_hub | |
| pip install huggingface_hub | |
| # Upload your LoRA | |
| from huggingface_hub import upload_file | |
| upload_file( | |
| path_or_fileobj="models/lora/your_lora.safetensors", | |
| path_in_repo="your_lora.safetensors", | |
| repo_id="your-username/your-lora-repo", | |
| repo_type="model" | |
| ) | |
| ``` | |
| 2. **Modify this Space to load your LoRA**: | |
| - Fork this Space or create a duplicate | |
| - Edit `backend/config.py`: | |
| ```python | |
| LORA_MODEL_NAME = "your_lora.safetensors" | |
| ``` | |
| - Add code in `app.py` to download from HF Hub: | |
| ```python | |
| from huggingface_hub import hf_hub_download | |
| lora_path = hf_hub_download( | |
| repo_id="your-username/your-lora-repo", | |
| filename="your_lora.safetensors", | |
| local_dir="models/lora" | |
| ) | |
| ``` | |
| 3. **Restart the Space** - Your custom LoRA will be loaded automatically! | |
| ### Option 2: Use Locally | |
| 1. **Place your LoRA file** in `models/lora/` directory | |
| 2. **Update config**: Set `LORA_MODEL_NAME` in `backend/config.py` | |
| 3. **Run backend**: `uvicorn main:app --reload` | |
| 4. The LoRA is automatically detected and loaded! | |
| ### How It Works | |
| When a LoRA is present: | |
| - ✅ Pipeline automatically loads it on startup | |
| - ✅ All enhancements use your custom style | |
| - ✅ No additional configuration needed | |
| - ✅ Can be combined with custom prompts for fine control | |
| ### Verify LoRA is Loaded | |
| Check the logs on startup: | |
| ``` | |
| INFO: Loading Stable Diffusion pipeline... | |
| INFO: LoRA model found: models/lora/your_lora.safetensors | |
| INFO: ✅ Custom LoRA loaded successfully! | |
| ``` | |
| ### Example: Architecture Firm Custom Style | |
| ```python | |
| # After training with your firm's rendering style | |
| # Place: models/lora/firm_style.safetensors | |
| # In config.py: | |
| LORA_MODEL_NAME = "firm_style.safetensors" | |
| # Now all enhancements will match your style! | |
| ``` | |
| --- | |
| """) | |
| gr.Markdown(""" | |
| ## 🚀 What is LoRA Training? | |
| LoRA (Low-Rank Adaptation) allows you to fine-tune the AI with just 10-50 image pairs to learn: | |
| - Your specific architectural rendering style | |
| - Preferred lighting and atmosphere | |
| - Consistent material treatments | |
| - Unique design aesthetics | |
| ### Benefits: | |
| - ✅ **Fast Training**: Only 1000 steps needed (~20-30 min on GPU) | |
| - ✅ **Small Models**: LoRA weights are only ~10-50 MB | |
| - ✅ **Style Consistency**: Perfect for architectural firms with specific styles | |
| - ✅ **Efficient**: Works on consumer GPUs | |
| --- | |
| ## 📋 Training Process Overview | |
| ### Step 1: Prepare Training Data | |
| Create **image pairs**: | |
| - **Input**: Your base architectural render (before) | |
| - **Target**: Your ideal enhanced result (after) | |
| **Requirements:** | |
| - Minimum: 10 pairs (recommended: 20-50) | |
| - Format: PNG or JPG | |
| - Resolution: 512x512 to 1024x1024 | |
| **Example structure:** | |
| ``` | |
| training_data/ | |
| inputs/ | |
| building_001_input.png | |
| building_002_input.png | |
| targets/ | |
| building_001_target.png | |
| building_002_target.png | |
| ``` | |
| ### Step 2: Setup Local Backend | |
| ```bash | |
| cd architecture-ai-enhancer/backend | |
| pip install -r requirements.txt | |
| uvicorn main:app --host 0.0.0.0 --port 8000 | |
| ``` | |
| ### Step 3: Upload Training Pairs | |
| Use the API at `http://localhost:8000/docs` to upload your image pairs: | |
| ```python | |
| import requests | |
| files = { | |
| 'input_image': open('building_001_input.png', 'rb'), | |
| 'target_image': open('building_001_target.png', 'rb') | |
| } | |
| response = requests.post( | |
| 'http://localhost:8000/training/upload_pair', | |
| files=files | |
| ) | |
| ``` | |
| ### Step 4: Start Training | |
| ```python | |
| config = { | |
| "train_steps": 1000, | |
| "learning_rate": 1e-4, | |
| "lora_rank": 8, | |
| "batch_size": 1 | |
| } | |
| response = requests.post( | |
| 'http://localhost:8000/training/start', | |
| json=config | |
| ) | |
| ``` | |
| ### Step 5: Use Your Custom Model | |
| Once trained, your custom LoRA is automatically used for all enhancements! | |
| --- | |
| ## 🎨 Training Tips | |
| **For Subtle Enhancements:** | |
| - train_steps: 500 | |
| - learning_rate: 5e-5 | |
| - lora_rank: 4 | |
| **For Dramatic Style Changes:** | |
| - train_steps: 1500 | |
| - learning_rate: 1e-4 | |
| - lora_rank: 12 | |
| **Balanced (Recommended):** | |
| - train_steps: 1000 | |
| - learning_rate: 1e-4 | |
| - lora_rank: 8 | |
| --- | |
| ## 📚 Complete Documentation | |
| For detailed step-by-step instructions, troubleshooting, and advanced techniques, see: | |
| **[📖 Complete LoRA Training Guide](https://huggingface.co/spaces/TransformacionDigitalAA/ArchEnhancer/blob/main/TRAINING_GUIDE.md)** | |
| This includes: | |
| - Detailed API usage examples | |
| - Complete Python training script | |
| - Troubleshooting common issues | |
| - Best practices for creating training data | |
| - How to deploy your custom LoRA | |
| --- | |
| ## 💡 Use Cases | |
| - **Architecture Firms**: Train on your signature rendering style | |
| - **Game Studios**: Consistent environmental concept art | |
| - **VFX Artists**: Specific lighting and atmosphere | |
| - **Real Estate**: Standardized visualization style | |
| --- | |
| ## 🔗 Resources | |
| - [Backend GitHub Repository](https://github.com/yourusername/architecture-ai-enhancer) | |
| - [LoRA Paper](https://arxiv.org/abs/2106.09685) | |
| - [Diffusers Documentation](https://huggingface.co/docs/diffusers/) | |
| """) | |
| # Tab 3: About | |
| with gr.Tab("ℹ️ About"): | |
| gr.Markdown(""" | |
| # About Architecture AI Enhancer | |
| ## 🔧 Technical Details | |
| - **Model**: Stable Diffusion 1.5 (runwayml/stable-diffusion-v1-5) | |
| - **Upscaler**: ESRGAN with fallback to Lanczos | |
| - **Framework**: PyTorch + Diffusers | |
| - **Interface**: Gradio 4.20.0 | |
| - **Version**: 1.0.0 | |
| ## ⚙️ Settings Guide | |
| ### Denoising Strength (0.1-0.8) | |
| Controls how much the AI modifies your input image: | |
| - **0.2-0.3**: Subtle enhancements (recommended for most cases) | |
| - **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 (recommended) | |
| - **7-10**: More stylized output | |
| - **11-15**: Very strong prompt adherence | |
| ### Custom Prompt | |
| Add specific details to guide the enhancement: | |
| - Lighting: "sunset lighting", "dramatic shadows" | |
| - Style: "modern minimalist", "brutalist concrete" | |
| - Materials: "glass facade", "wooden accents" | |
| - Atmosphere: "foggy morning", "golden hour" | |
| ## 🚀 Features | |
| - ✅ AI-powered enhancement with Stable Diffusion | |
| - ✅ 2x image upscaling | |
| - ✅ Professional post-processing | |
| - ✅ Custom prompt support | |
| - ✅ Adjustable parameters | |
| - ✅ Custom LoRA training (local) | |
| ## 📊 Performance | |
| - **CPU (HF Spaces Free Tier)**: 2-5 minutes per image | |
| - **GPU (Local/Paid)**: ~30 seconds per image | |
| ## 🙏 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) | |
| ## 📝 License | |
| MIT License - Free for commercial and personal use | |
| --- | |
| **Made with ❤️ for the architecture community** | |
| """) | |
| if __name__ == "__main__": | |
| # Check if running on HF Spaces | |
| is_spaces = os.getenv("SPACE_ID") is not None | |
| if is_spaces: | |
| # HF Spaces specific configuration | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False | |
| ) | |
| else: | |
| # Local development | |
| demo.launch(share=True) | |