upload 3 files

README.md

@@ -1,13 +1,123 @@
 ---
-title: Controlnet Floorplan Demo
-emoji: 🌖
-colorFrom: indigo
+title: Draw Your Floorplan - ControlNet
+emoji: 🏠
+colorFrom: blue
 colorTo: green
 sdk: gradio
-sdk_version: 6.2.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 license: mit
+models:
+  - Qistinasofea/controlnet-floorplan
+  - stable-diffusion-v1-5/stable-diffusion-v1-5
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🏠 Draw Your Floorplan - ControlNet
+
+**AI54 Final Project - Spatially Conditioned Floorplan Generation**
+
+## 🎨 Interactive Demo
+
+This Space allows you to **draw colored segmentation masks** and generate architectural floorplans using a fine-tuned ControlNet model.
+
+### How to Use:
+
+1. **Draw** colored regions on the canvas - each color represents a different room type
+2. **Describe** your floorplan in the text box
+3. **Adjust** settings if needed (inference steps, control strength, seed)
+4. Click **Generate Floorplan** to see your AI-generated layout!
+
+### Suggested Colors:
+- 🔴 Red - Living room / Main spaces
+- 🟢 Green - Bedrooms
+- 🔵 Blue - Bathrooms
+- 🟡 Yellow - Kitchen
+- 🟣 Purple - Dining area
+- 🟠 Orange - Office / Study
+- 🩵 Cyan - Utility / Storage
+
+---
+
+## 📊 Model Information
+
+### Training Details:
+- **Method:** Full ControlNet Fine-Tuning
+- **Base Model:** Stable Diffusion 1.5 (frozen)
+- **ControlNet:** Segmentation variant (fully trained)
+- **Dataset:** 11,375 orientation-normalized floorplan samples
+- **Parameters:** 361M trainable parameters (100% of ControlNet)
+- **Training Steps:** 10,000
+- **Final Loss:** 0.0887
+- **Training Time:** 3.7 hours on T4 GPU
+
+### Architecture:
+The model uses a two-stage architecture:
+1. **Base Model (SD 1.5):** Generates realistic textures and appearance (frozen weights)
+2. **ControlNet:** Guides spatial structure based on colored segmentation input (fully fine-tuned)
+
+This separation allows the model to:
+- ✅ Preserve spatial layouts from user drawings
+- ✅ Generate realistic architectural details
+- ✅ Maintain consistent room boundaries
+- ✅ Produce diverse outputs from the same layout
+
+---
+
+## 🔗 Links
+
+- **Trained Model:** [Qistinasofea/controlnet-floorplan](https://huggingface.co/Qistinasofea/controlnet-floorplan)
+- **Dataset:** [Qistinasofea/floorplan-12k-aligned](https://huggingface.co/datasets/Qistinasofea/floorplan-12k-aligned)
+- **Training Notebook:** Available in model repository
+
+---
+
+## 🎓 Academic Context
+
+This is a final project for **AI54: Artificial Intelligence course** focused on:
+- Conditional image generation
+- Spatial control in diffusion models
+- ControlNet architecture and training
+- Parameter-efficient fine-tuning considerations
+- Real-world application development
+
+### Key Contributions:
+1. **Dataset Preprocessing:** Orientation normalization using PCA-based rotation alignment
+2. **Training Strategy:** Full fine-tuning justified by dataset size (11,375 samples)
+3. **User Interface:** Visual layout-driven interaction for non-technical users
+
+---
+
+## 💻 Technical Stack
+
+- **Framework:** 🤗 Diffusers
+- **Model:** ControlNet + Stable Diffusion 1.5
+- **Interface:** Gradio
+- **Deployment:** HuggingFace Spaces
+- **Hardware:** GPU-enabled (T4 or better recommended)
+
+---
+
+## 📝 Citation
+
+If you use this model or approach in your work, please cite:
+
+```
+@misc{controlnet-floorplan-2024,
+  author = {Qistinasofea},
+  title = {ControlNet for Floorplan Generation},
+  year = {2024},
+  publisher = {HuggingFace},
+  howpublished = {\url{https://huggingface.co/Qistinasofea/controlnet-floorplan}}
+}
+```
+
+---
+
+## 📄 License
+
+This project is released under the MIT License. The base Stable Diffusion 1.5 model follows its original CreativeML Open RAIL-M license.
+
+---
+
+**Built with ❤️ for AI54 Final Project**

app.py

@@ -0,0 +1,289 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from PIL import Image
+import numpy as np
+
+# ============================================================================
+# MODEL LOADING
+# ============================================================================
+
+print("🔄 Loading ControlNet Floorplan model...")
+
+# Load trained ControlNet
+controlnet = ControlNetModel.from_pretrained(
+    "Qistinasofea/controlnet-floorplan",
+    torch_dtype=torch.float16
+)
+
+# Create pipeline
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    torch_dtype=torch.float16,
+    safety_checker=None
+)
+
+# Use faster scheduler
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+
+# Move to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+
+print(f"✅ Model loaded on {device}!")
+
+# ============================================================================
+# GENERATION FUNCTION
+# ============================================================================
+
+def generate_floorplan(
+    segmentation_mask,
+    prompt,
+    num_steps,
+    conditioning_scale,
+    seed
+):
+    """
+    Generate floorplan from colored segmentation mask
+    
+    Args:
+        segmentation_mask: User-drawn colored layout
+        prompt: Text description
+        num_steps: Number of diffusion steps
+        conditioning_scale: ControlNet influence strength
+        seed: Random seed for reproducibility
+    """
+    
+    if segmentation_mask is None:
+        return None, "⚠️ Please draw a layout first!"
+    
+    try:
+        # Handle Gradio ImageEditor format
+        if isinstance(segmentation_mask, dict):
+            # Extract the composite image
+            if 'composite' in segmentation_mask:
+                segmentation_mask = segmentation_mask['composite']
+            elif 'background' in segmentation_mask:
+                segmentation_mask = segmentation_mask['background']
+        
+        # Convert to PIL Image
+        if isinstance(segmentation_mask, np.ndarray):
+            segmentation_mask = Image.fromarray(segmentation_mask)
+        
+        # Ensure RGB and correct size
+        segmentation_mask = segmentation_mask.convert('RGB')
+        segmentation_mask = segmentation_mask.resize((512, 512), Image.LANCZOS)
+        
+        # Generate floorplan
+        generator = torch.Generator(device=device).manual_seed(int(seed))
+        
+        output = pipe(
+            prompt=prompt,
+            image=segmentation_mask,
+            num_inference_steps=int(num_steps),
+            controlnet_conditioning_scale=float(conditioning_scale),
+            generator=generator
+        ).images[0]
+        
+        status = f"✅ Generated successfully! (seed: {seed})"
+        return output, status
+        
+    except Exception as e:
+        error_msg = f"❌ Error: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+
+# ============================================================================
+# GRADIO INTERFACE
+# ============================================================================
+
+# Custom CSS for better styling
+custom_css = """
+#title {
+    text-align: center;
+    font-size: 2.5em;
+    font-weight: bold;
+    margin-bottom: 10px;
+}
+#subtitle {
+    text-align: center;
+    font-size: 1.2em;
+    color: #666;
+    margin-bottom: 20px;
+}
+.color-legend {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 10px;
+    margin: 10px 0;
+}
+.color-item {
+    display: flex;
+    align-items: center;
+    gap: 5px;
+}
+.color-box {
+    width: 20px;
+    height: 20px;
+    border: 1px solid #ccc;
+}
+"""
+
+# Create interface
+with gr.Blocks(css=custom_css, title="Draw Your Floorplan") as demo:
+    
+    # Header
+    gr.HTML("<h1 id='title'>🏠 Draw Your Floorplan</h1>")
+    gr.HTML("<p id='subtitle'>AI54 Final Project - ControlNet for Architectural Layout Generation</p>")
+    
+    gr.Markdown("""
+    ### 📝 How to Use:
+    1. **Draw** colored regions on the canvas - each color represents a room type
+    2. **Describe** your floorplan in the text box
+    3. **Adjust** settings if needed (optional)
+    4. Click **Generate Floorplan** 🎨
+    
+    ### 🎨 Suggested Colors:
+    - 🔴 **Red** - Living room / Main spaces
+    - 🟢 **Green** - Bedrooms
+    - 🔵 **Blue** - Bathrooms
+    - 🟡 **Yellow** - Kitchen
+    - 🟣 **Purple** - Dining area
+    - 🟠 **Orange** - Office / Study
+    - 🩵 **Cyan** - Utility / Storage
+    - 🟤 **Brown** - Corridors / Hallways
+    """)
+    
+    with gr.Row():
+        # LEFT COLUMN - Input
+        with gr.Column(scale=1):
+            gr.Markdown("### 🎨 Step 1: Draw Your Layout")
+            
+            # Drawing canvas
+            input_image = gr.ImageEditor(
+                label="Draw colored regions (each color = room type)",
+                type="numpy",
+                image_mode="RGB",
+                brush=gr.Brush(
+                    colors=["#FF0000", "#00FF00", "#0000FF", "#FFFF00", 
+                           "#FF00FF", "#FFA500", "#00FFFF", "#8B4513"],
+                    default_size=20
+                ),
+                height=512,
+                width=512
+            )
+            
+            gr.Markdown("### 📝 Step 2: Describe Your Floorplan")
+            
+            prompt = gr.Textbox(
+                label="Text Description",
+                placeholder="Example: A modern apartment with 2 bedrooms, kitchen, and living room",
+                value="A residential floorplan with multiple rooms",
+                lines=3
+            )
+            
+            # Advanced settings
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                num_steps = gr.Slider(
+                    minimum=10,
+                    maximum=50,
+                    value=20,
+                    step=1,
+                    label="Inference Steps (higher = better quality, slower)"
+                )
+                
+                conditioning_scale = gr.Slider(
+                    minimum=0.5,
+                    maximum=2.0,
+                    value=1.0,
+                    step=0.1,
+                    label="Control Strength (how closely to follow your drawing)"
+                )
+                
+                seed = gr.Slider(
+                    minimum=0,
+                    maximum=999999,
+                    value=42,
+                    step=1,
+                    label="Random Seed (for reproducibility)"
+                )
+            
+            # Generate button
+            generate_btn = gr.Button("🎨 Generate Floorplan", variant="primary", size="lg")
+        
+        # RIGHT COLUMN - Output
+        with gr.Column(scale=1):
+            gr.Markdown("### 🖼️ Generated Floorplan")
+            
+            output_image = gr.Image(
+                label="Your AI-Generated Floorplan",
+                type="pil",
+                height=512
+            )
+            
+            status_text = gr.Textbox(
+                label="Status",
+                value="Draw a layout and click Generate!",
+                interactive=False
+            )
+            
+            gr.Markdown("""
+            ### 💡 Tips for Best Results:
+            - Draw clear, distinct colored regions
+            - Use different colors for different room types
+            - Make rooms rectangular for better results
+            - Add a descriptive text prompt
+            - Experiment with different seeds for variety!
+            
+            ### 📊 Model Info:
+            - **Training:** Full ControlNet fine-tuning on 11,375 samples
+            - **Architecture:** SD 1.5 + ControlNet Segmentation
+            - **Parameters:** 361M trainable parameters
+            - **Final Loss:** 0.0887
+            """)
+    
+    # Examples section
+    gr.Markdown("---")
+    gr.Markdown("### 📋 Example Layouts")
+    gr.Markdown("Click an example below to try it out!")
+    
+    # Note: You'll need to add actual example images
+    # For now, using placeholders
+    gr.Examples(
+        examples=[
+            [None, "A small studio apartment with combined living and sleeping area", 20, 1.0, 42],
+            [None, "A two-bedroom apartment with separate kitchen and bathroom", 20, 1.0, 123],
+            [None, "An office space with multiple workrooms and meeting area", 20, 1.0, 456],
+        ],
+        inputs=[input_image, prompt, num_steps, conditioning_scale, seed],
+        outputs=[output_image, status_text],
+        fn=generate_floorplan,
+        cache_examples=False,
+    )
+    
+    # Footer
+    gr.Markdown("---")
+    gr.Markdown("""
+    ### 🎓 AI54 Final Project
+    **Student:** Qistinasofea  
+    **Model:** [Qistinasofea/controlnet-floorplan](https://huggingface.co/Qistinasofea/controlnet-floorplan)  
+    **Dataset:** [Qistinasofea/floorplan-12k-aligned](https://huggingface.co/datasets/Qistinasofea/floorplan-12k-aligned)
+    
+    Built with 🤗 Diffusers, Gradio, and ControlNet
+    """)
+    
+    # Connect button to function
+    generate_btn.click(
+        fn=generate_floorplan,
+        inputs=[input_image, prompt, num_steps, conditioning_scale, seed],
+        outputs=[output_image, status_text]
+    )
+
+# ============================================================================
+# LAUNCH
+# ============================================================================
+
+if __name__ == "__main__":
+    demo.queue()  # Enable queuing for better performance
+    demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+diffusers==0.30.0
+transformers
+accelerate
+gradio
+Pillow
+numpy
+safetensors