working demo, organized files

gradio_app/.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

gradio_app/app.py

@@ -0,0 +1,1504 @@
+import os
+import os.path as osp 
+import sys
+import numpy as np
+import tempfile
+import shutil 
+import base64
+import io
+from PIL import Image
+import gradio as gr
+import time
+import copy
+import requests
+import json
+import pickle 
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from object_scales import scales  
+from transformers import CLIPTokenizer, PretrainedConfig, T5TokenizerFast
+import pickle 
+from datetime import datetime 
+from infer_backend import initialize_inference_engine, run_inference_from_gradio
+
+COLORS = [
+    (1.0, 0.0, 0.0),    # Red
+    (0.0, 0.8, 0.2),    # Green
+    (0.0, 0.0, 1.0),    # Blue
+    (1.0, 1.0, 0.0),    # Yellow
+    (0.0, 1.0, 1.0),    # Cyan
+    (1.0, 0.0, 1.0),    # Magenta
+    (1.0, 0.6, 0.0),    # Orange
+    (0.6, 0.0, 0.8),    # Purple
+    (0.0, 0.4, 0.0),    # Dark Green
+    (0.8, 0.8, 0.8),    # Light Gray
+    (0.2, 0.2, 0.2)     # Dark Gray
+]
+
+CHECKPOINT_NAMES = [
+    "rgb__r1/epoch-0__checkpoint-25917",
+    "rgb__finetune_1024/epoch-0__checkpoint-3000", 
+    "rgb__finetune_1024/epoch-1__checkpoint-4000", 
+    "rgb__finetune_1024/epoch-1__checkpoint-5000", 
+    "rgb__finetune_1024/epoch-1__checkpoint-6000", 
+    "rgb__finetune_1024/epoch-1__checkpoint-7000", 
+    "rgb__finetune_1024/epoch-1__checkpoint-7932", 
+]
+
+PRETRAINED_MODEL_NAME_OR_PATH = "black-forest-labs/FLUX.1-dev" 
+
+tokenizer = T5TokenizerFast.from_pretrained(
+    PRETRAINED_MODEL_NAME_OR_PATH, 
+    subfolder="tokenizer_2",
+    revision=None, 
+)
+
+placeholder_token_str = ["<placeholder>"]
+num_added_tokens = tokenizer.add_tokens(placeholder_token_str)   
+assert num_added_tokens == 1 
+
+def generate_image_event(camera_elevation, camera_lens, surrounding_prompt, checkpoint_name, 
+                        height, width, seed, guidance_scale, num_steps):
+    """Generate final image with segmentation masks and run inference"""
+    # Update scene manager's inference params before generation
+    scene_manager.update_inference_params(height, width, seed, guidance_scale, num_steps, checkpoint_name)
+    if not scene_manager.objects:
+        return (
+            "⚠️ No objects to render",
+            gr.update(),
+            Image.new('RGB', (512, 512), color='white')
+        )
+    
+    # Get subject descriptions
+    subject_descriptions = [obj['description'] for obj in scene_manager.objects]
+    
+    print(f"Surrounding prompt: {surrounding_prompt}")
+    print(f"Subject descriptions: {subject_descriptions}")
+    print(f"Selected checkpoint: {checkpoint_name}")
+
+    placeholder_prompt = "a photo of PLACEHOLDER " + surrounding_prompt 
+
+    # Create placeholder text
+    subject_embeds = [] 
+    for subject_idx, subject_desc in enumerate(subject_descriptions): 
+        input_ids = tokenizer.encode(subject_desc, return_tensors="pt", max_length=77)[0] 
+        subject_embed = {"input_ids_t5": input_ids.tolist()} 
+        subject_embeds.append(subject_embed)
+
+    placeholder_text = ""
+    for subject in subject_descriptions[:-1]:
+        placeholder_text = placeholder_text + f"<placeholder> {subject} and "
+    for subject in subject_descriptions[-1:]:
+        placeholder_text = placeholder_text + f"<placeholder> {subject}"
+    placeholder_text = placeholder_text.strip()
+
+    placeholder_token_prompt = placeholder_prompt.replace("PLACEHOLDER", placeholder_text) 
+
+    call_ids = get_call_ids_from_placeholder_prompt_flux(prompt=placeholder_token_prompt, 
+                                             subjects=subject_descriptions, 
+                                             subjects_embeds=subject_embeds,
+                                             debug=True
+                                            ) 
+    print(f"Generated call IDs: {call_ids}") 
+
+    # Convert to server expected format
+    subjects_data, camera_data = scene_manager._convert_to_blender_format()
+    
+    # Render final high-quality image using CYCLES (port 5002)
+    final_img = scene_manager.render_client._send_render_request(
+        scene_manager.render_client.final_server_url, 
+        subjects_data, 
+        camera_data
+    )
+
+    final_img.save("model_condition.jpg") 
+    
+    # Render segmentation masks
+    success, segmask_images, error_msg = scene_manager.render_client.render_segmasks(subjects_data, camera_data)
+    
+    if not success:
+        return (
+            f"❌ Failed to render segmentation masks: {error_msg}",
+            gr.update(),
+            Image.new('RGB', (512, 512), color='white')
+        )
+
+    # Save all files to the correct location  
+    root_save_dir = "/archive/vaibhav.agrawal/a-bev-of-the-latents/gradio_files/" 
+    os.system(f"rm -f {root_save_dir}/*") 
+    
+    # Save final render to root directory
+    final_render_path = osp.join(root_save_dir, "cv_render.jpg")
+    final_img.save(final_render_path)
+    
+    # Move segmentation masks
+    for subject_idx in range(len(subject_descriptions)): 
+        shutil.move(
+            f"{str(subject_idx).zfill(3)}_segmask_cv.png", 
+            osp.join(root_save_dir, f"main__segmask_{str(subject_idx).zfill(3)}__{1.00}.png")
+        ) 
+    
+    # Create JSONL
+    jsonl = [{
+        "cv": final_render_path,
+        "target": final_render_path,
+        "cuboids_segmasks": [
+            osp.join(root_save_dir, f"main__segmask_{str(subject_idx).zfill(3)}__{1.00}.png") 
+            for subject_idx in range(len(subject_descriptions))
+        ],
+        "PLACEHOLDER_prompts": placeholder_prompt, 
+        "subjects": subject_descriptions, 
+        "call_ids": call_ids, 
+    }] 
+
+    jsonl_path = osp.join(root_save_dir, "cuboids.jsonl")
+    with open(jsonl_path, "w") as f: 
+        json.dump(jsonl[0], f)
+    
+    # Run inference using the pre-loaded model
+    print(f"\n{'='*60}")
+    print(f"RUNNING INFERENCE")
+    print(f"{'='*60}\n")
+    
+    inference_success, generated_image, inference_msg = run_inference_from_gradio(
+        checkpoint_name=checkpoint_name,
+        height=height,
+        width=width,
+        seed=seed,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_steps,
+        jsonl_path=jsonl_path
+    )
+    
+    if not inference_success:
+        return (
+            f"✅ Saved files but inference failed: {inference_msg}",
+            final_img,
+            Image.new('RGB', (512, 512), color='white')
+        )
+    
+    status_msg = f"✅ Generated image using {checkpoint_name} with {len(segmask_images)} segmentation masks"
+
+    # Render final high-quality image using CYCLES (port 5002)
+    final_img = scene_manager.render_client._send_render_request(
+        scene_manager.render_client.paper_figure_server_url, 
+        subjects_data, 
+        camera_data
+    )
+    
+    return (
+        status_msg,
+        final_img,  # Display CV render in Camera View
+        generated_image  # Display generated image in Generated Image section
+    )
+
+
+def get_call_ids_from_placeholder_prompt_flux(prompt: str, subjects, subjects_embeds: list, debug: bool):  
+	assert prompt.find("<placeholder>") != -1, "Prompt must contain <placeholder> to get call ids" 
+
+	# the placeholder token ID for all the tokenizers 
+	placeholder_token_three = tokenizer.encode("<placeholder>", return_tensors="pt")[0][:-1].item()  
+	prompt_tokens_three = tokenizer.encode(prompt, return_tensors="pt")[0].tolist()  
+
+	placeholder_token_locations_three = [i for i, w in enumerate(prompt_tokens_three) if w == placeholder_token_three] 
+	prompt = prompt.replace("<placeholder> ", "") 
+
+
+	call_ids = [] 
+	for subject_idx, (subject, subject_embed) in enumerate(zip(subjects, subjects_embeds)):  
+		subject_prompt_ids_t5 = subject_embed["input_ids_t5"][:-1] # T5 has SOT token only  
+		num_t5_tokens_subject = len(subject_prompt_ids_t5) 
+
+		t5_call_ids_subject = [i + placeholder_token_locations_three[subject_idx] - 2 * subject_idx - 1 for i in range(num_t5_tokens_subject)] 
+		call_ids.append(t5_call_ids_subject) 
+
+		prompt_wo_placeholder = prompt.replace("<placeholder> ", "") 
+		t5_call_strs = tokenizer.batch_decode(tokenizer.encode(prompt_wo_placeholder, return_tensors="pt")[0].tolist())  
+		t5_call_strs = [t5_call_strs[i] for i in t5_call_ids_subject] 
+		if debug: 
+			print(f"{prompt = }, t5 CALL strs for {subject} = {t5_call_strs}") 
+
+	return call_ids  
+
+
+def map_point_to_rgb(x, y):
+    """
+    Map (x, y) inside the frustum to an RGB color with continuity and variation.
+    """
+    # Frustum boundaries
+    X_MIN, X_MAX = -10.0, -1.0
+    Y_MIN_AT_XMIN, Y_MAX_AT_XMIN = -4.5, 4.5
+    Y_MIN_AT_XMAX, Y_MAX_AT_XMAX = -0.5, 0.5
+    
+    # Normalize x to [0, 1]
+    x_norm = (x - X_MIN) / (X_MAX - X_MIN)
+    # x_norm = np.clip(x_norm, 0, 1)
+
+    # Compute current Y bounds at given x using linear interpolation
+    y_min = Y_MIN_AT_XMIN + x_norm * (Y_MIN_AT_XMAX - Y_MIN_AT_XMIN)
+    y_max = Y_MAX_AT_XMIN + x_norm * (Y_MAX_AT_XMAX - Y_MAX_AT_XMIN)
+
+    # Normalize y to [0, 1] within current bounds
+    if y_max != y_min:
+        y_norm = (y - y_min) / (y_max - y_min)
+    else:
+        y_norm = 0.5
+    y_norm = np.clip(y_norm, 0.0, 1.0)
+
+    # Color mapping: more variation along x
+    r = x_norm
+    g = y_norm 
+    b = 1.0 - x_norm
+
+    return (r, g, b)
+
+def rgb_to_hex(rgb_tuple):
+    """Convert RGB tuple (0-1 range) to hex color string."""
+    r, g, b = rgb_tuple
+    return f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}"
+
+
+class BlenderRenderClient:
+    def __init__(self, cv_server_url="http://127.0.0.1:5001", segmask_server_url="http://127.0.0.1:5003", final_server_url="http://127.0.0.1:5002", paper_figure_server_url="http://127.0.0.1:5004"): 
+        """
+        Initialize the Blender render client.
+        
+        Args:
+            cv_server_url (str): URL of the camera view render server
+            segmask_server_url (str): URL of the segmentation mask render server
+        """
+        self.cv_server_url = cv_server_url
+        self.segmask_server_url = segmask_server_url
+        self.final_server_url = final_server_url 
+        self.paper_figure_server_url = paper_figure_server_url 
+        self.timeout = 30  # 30 second timeout for renders
+
+    def render_segmasks(self, subjects_data: list, camera_data: dict) -> tuple:
+        """
+        Send a segmentation mask render request. 
+        Returns (success: bool, segmask_images: list of PIL Images or None, error_message: str or None)
+        """
+        try:
+            request_data = {
+                "subjects_data": subjects_data,
+                "camera_data": camera_data,
+                "num_samples": 1
+            }
+            
+            response = requests.post(
+                f"{self.segmask_server_url}/render_segmasks",
+                json=request_data,
+                timeout=self.timeout
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result["success"]:
+                    # Decode all segmentation masks
+                    segmask_images = []
+                    for img_base64 in result["segmasks_base64"]:
+                        img_data = base64.b64decode(img_base64)
+                        img = Image.open(io.BytesIO(img_data))
+                        segmask_images.append(img)
+                    
+                    print(f"Successfully rendered {len(segmask_images)} segmentation masks")
+                    return True, segmask_images, None
+                else:
+                    error_msg = result.get('error_message', 'Unknown error')
+                    print(f"Segmask render failed: {error_msg}")
+                    return False, None, error_msg
+            else:
+                error_msg = f"HTTP error {response.status_code}: {response.text}"
+                print(error_msg)
+                return False, None, error_msg
+                
+        except requests.exceptions.Timeout:
+            error_msg = "Segmask render request timed out"
+            print(error_msg)
+            return False, None, error_msg
+        except Exception as e:
+            error_msg = f"Segmask render request failed: {e}"
+            print(error_msg)
+            return False, None, error_msg
+
+        
+    def _send_render_request(self, server_url: str, subjects_data: list, camera_data: dict) -> Image.Image:
+        """Send a render request to a server and return the image."""
+        try:
+            request_data = {
+                "subjects_data": subjects_data,
+                "camera_data": camera_data,
+                "num_samples": 1
+            }
+            print(f"passing {subjects_data = } to server at {server_url}") 
+            
+            response = requests.post(
+                f"{server_url}/render",
+                json=request_data,
+                timeout=self.timeout
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result["success"]:
+                    # Decode base64 image
+                    img_data = base64.b64decode(result["image_base64"])
+                    img = Image.open(io.BytesIO(img_data))
+                    return img
+                else:
+                    print(f"Render failed: {result.get('error_message', 'Unknown error')}")
+                    return self._create_error_image("red")
+            else:
+                print(f"HTTP error {response.status_code}: {response.text}")
+                return self._create_error_image("orange")
+                
+        except requests.exceptions.Timeout:
+            print("Render request timed out")
+            return self._create_error_image("yellow")
+        except Exception as e:
+            print(f"Render request failed: {e}")
+            return self._create_error_image("red")
+    
+    def _create_error_image(self, color: str) -> Image.Image:
+        """Create a colored error image."""
+        return Image.new('RGB', (512, 512), color=color)
+    
+# --- Scene Management Class ---
+class SceneManager:
+    def __init__(self):
+        self.objects = []
+        self.camera_elevation = 30.0
+        self.camera_lens = 50.0
+        self.surrounding_prompt = ""
+        self.next_color_idx = 0
+        self.colors = [
+            (1.0, 0.0, 0.0),  # red
+            (0.0, 0.0, 1.0),  # blue
+            (0.0, 1.0, 0.0),  # green
+            (0.5, 0.0, 0.5),  # purple
+            (1.0, 0.5, 0.0),  # orange
+            (1.0, 1.0, 0.0),  # yellow
+            (0.0, 1.0, 1.0),  # cyan
+            (1.0, 0.0, 1.0),  # magenta
+        ]
+        
+        # Add inference parameters with defaults
+        self.inference_params = {
+            'height': 512,
+            'width': 512,
+            'seed': 42,
+            'guidance_scale': 3.5,
+            'num_inference_steps': 25,
+            'checkpoint': CHECKPOINT_NAMES[0] if CHECKPOINT_NAMES else None
+        }
+        
+        # Initialize BlenderRenderClient
+        self.render_client = BlenderRenderClient()
+        
+        # Load asset dimensions
+        self.asset_dimensions = self._load_asset_dimensions()
+
+
+    def update_inference_params(self, height, width, seed, guidance_scale, num_steps, checkpoint):
+        """Update inference parameters"""
+        self.inference_params = {
+            'height': height,
+            'width': width,
+            'seed': seed,
+            'guidance_scale': guidance_scale,
+            'num_inference_steps': num_steps,
+            'checkpoint': checkpoint
+        }
+
+
+    def update_cuboid_description(self, obj_id, new_description):
+        """Update the description of a cuboid"""
+        if 0 <= obj_id < len(self.objects):
+            if new_description.strip():  # Check not empty
+                self.objects[obj_id]['description'] = new_description.strip()
+                return True
+        return False
+
+
+    def save_scene_to_pkl(self, filepath=None):
+        """Save current scene data to pkl file including inference parameters"""
+        if filepath is None:
+            # Auto-generate filename with timestamp
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filepath = f"scene_{timestamp}.pkl"
+        
+        # Convert to the expected format
+        subjects_data = []
+        for obj in self.objects:
+            subject_dict = {
+                'name': obj['description'],
+                'type': obj['type'],  # Save the object type
+                'dims': tuple(obj['size']),  # (width, depth, height)
+                'x': [obj['position'][0] - 6.0],
+                'y': [obj['position'][1]],
+                'z': [obj['position'][2]],
+                'azimuth': [np.radians(obj['azimuth'])],  # Convert to radians
+                'bbox': [(0, 0, 0, 0)]  # Placeholder, can be computed if needed
+            }
+            subjects_data.append(subject_dict)
+        
+        camera_data = {
+            'camera_elevation': np.radians(self.camera_elevation),
+            'lens': self.camera_lens,
+            'global_scale': 1.0  # Default value
+        }
+        
+        scene_dict = {
+            'subjects_data': subjects_data,
+            'camera_data': camera_data,
+            'surrounding_prompt': self.surrounding_prompt,
+            'inference_params': self.inference_params.copy()
+        }
+        
+        try:
+            with open(filepath, 'wb') as f:
+                pickle.dump(scene_dict, f)
+            return True, filepath, None
+        except Exception as e:
+            return False, None, str(e)
+
+
+    def load_scene_from_pkl(self, filepath):
+        """Load scene data from pkl file including inference parameters"""
+        try:
+            with open(filepath, 'rb') as f:
+                scene_dict = pickle.load(f)
+            
+            # Clear existing objects
+            self.objects = []
+            self.next_color_idx = 0
+            
+            # Load subjects
+            subjects_data = scene_dict.get('subjects_data', [])
+            for subject_dict in subjects_data:
+                name = subject_dict.get('name', 'Loaded Object')
+                asset_type = subject_dict.get('type', 'Custom')  # Load the type
+                dims = subject_dict.get('dims', (1.0, 1.0, 1.0))
+                x = float(subject_dict.get('x', [0.0])[0]) + 6.0  
+                y = float(subject_dict.get('y', [0.0])[0])
+                z = float(subject_dict.get('z', [0.0])[0]) 
+                azimuth_rad = float(subject_dict.get('azimuth', [0.0])[0])
+                azimuth_deg = np.degrees(azimuth_rad)
+                
+                # Determine original_asset_size based on type
+                if asset_type == "Custom" or asset_type not in self.asset_dimensions:
+                    original_asset_size = None
+                else:
+                    # Look up the original asset dimensions
+                    asset_dims = self.asset_dimensions[asset_type]
+                    original_asset_size = [float(asset_dims[0]), float(asset_dims[1]), float(asset_dims[2])]
+                
+                # Create object
+                obj_id = len(self.objects)
+                size_list = [float(d) for d in dims]
+                cuboid = {
+                    'id': obj_id,
+                    'description': name,
+                    'type': asset_type,  # Use the loaded type
+                    'position': [x, y, z],
+                    'size': size_list,
+                    'original_asset_size': original_asset_size,  # Restore from asset_dimensions
+                    'azimuth': float(azimuth_deg),
+                    'color': self._get_next_color()
+                }
+                self.objects.append(cuboid)
+            
+            # Load camera settings
+            camera_data = scene_dict.get('camera_data', {})
+            camera_elev_rad = float(camera_data.get('camera_elevation', np.radians(30.0)))
+            self.camera_elevation = float(np.degrees(camera_elev_rad))
+            self.camera_lens = float(camera_data.get('lens', 50.0))
+            
+            # Load surrounding prompt
+            self.surrounding_prompt = scene_dict.get('surrounding_prompt', '')
+            
+            # Load inference parameters
+            loaded_inference_params = scene_dict.get('inference_params', {})
+            
+            # Get checkpoint, fall back to first available if not found
+            saved_checkpoint = loaded_inference_params.get('checkpoint')
+            if saved_checkpoint and saved_checkpoint in CHECKPOINT_NAMES:
+                checkpoint = saved_checkpoint
+            else:
+                checkpoint = CHECKPOINT_NAMES[0] if CHECKPOINT_NAMES else None
+                if saved_checkpoint:
+                    print(f"Warning: Saved checkpoint '{saved_checkpoint}' not found, using '{checkpoint}' instead")
+            
+            self.inference_params = {
+                'height': loaded_inference_params.get('height', 512),
+                'width': loaded_inference_params.get('width', 512),
+                'seed': loaded_inference_params.get('seed', 42),
+                'guidance_scale': loaded_inference_params.get('guidance_scale', 3.5),
+                'num_inference_steps': loaded_inference_params.get('num_inference_steps', 25),
+                'checkpoint': checkpoint
+            }
+            
+            return True, len(subjects_data), None
+        except FileNotFoundError:
+            return False, 0, f"File not found: {filepath}"
+        except Exception as e:
+            return False, 0, f"Error loading file: {str(e)}"
+
+
+    def _load_asset_dimensions(self):
+        """Load asset dimensions from pickle file"""
+        pkl_path = "asset_dimensions.pkl"
+        if os.path.exists(pkl_path):
+            try:
+                with open(pkl_path, 'rb') as f:
+                    return pickle.load(f)
+            except Exception as e:
+                print(f"Warning: Could not load asset dimensions: {e}")
+                return {}
+        else:
+            print(f"Warning: asset_dimensions.pkl not found at {pkl_path}")
+            return {}
+
+    def get_asset_type_choices(self):
+        """Get list of asset types for dropdown"""
+        choices = ["Custom"]
+        if self.asset_dimensions:
+            choices.extend(sorted(self.asset_dimensions.keys()))
+        return choices
+
+    def _get_next_color(self):
+        color = self.colors[self.next_color_idx % len(self.colors)]
+        self.next_color_idx += 1
+        return color
+
+
+    def harmonize_scales(self):
+        """
+        Harmonize the scales of all non-Custom objects based on object scales.
+        Always scales from original asset dimensions, ignoring any manual edits.
+        Custom objects remain unchanged.
+        """
+        if not self.objects:
+            return "No objects to harmonize"
+        
+        # Find objects that can be harmonized (non-Custom with valid scales and original_asset_size)
+        harmonizable_objects = []
+        for obj in self.objects:
+            if (obj['type'] != "Custom" and 
+                obj['type'] in scales and 
+                obj['original_asset_size'] is not None):
+                harmonizable_objects.append(obj)
+        
+        if not harmonizable_objects:
+            return "No objects with defined scales to harmonize (all are Custom)"
+        
+        # Find the largest scale among harmonizable objects
+        max_scale = max(scales[obj['type']] for obj in harmonizable_objects)
+        
+        if max_scale == 0:
+            return "Invalid max scale (0)"
+        
+        # Harmonize each object by scaling from ORIGINAL ASSET dimensions
+        for obj in harmonizable_objects:
+            obj_scale = scales[obj['type']]
+            scale_factor = obj_scale / max_scale
+            
+            # Scale from ORIGINAL ASSET dimensions, not current dimensions
+            obj['size'][0] = obj['original_asset_size'][0] * scale_factor  # width
+            obj['size'][1] = obj['original_asset_size'][1] * scale_factor  # depth
+            obj['size'][2] = obj['original_asset_size'][2] * scale_factor  # height
+            
+            # Update z position to keep object on ground
+            obj['position'][2] = 0.0 
+        
+        return f"Harmonized {len(harmonizable_objects)} objects based on largest scale: {max_scale}"
+
+
+    def add_cuboid(self, description="New Cuboid", asset_type="Custom"):
+        """Add a cuboid with dimensions based on asset type"""
+        obj_id = len(self.objects)
+        
+        # Determine dimensions based on asset type
+        if asset_type == "Custom" or asset_type not in self.asset_dimensions:
+            size = [1.0, 1.0, 1.0]  # Default size
+            original_asset_size = None  # Custom objects have no original asset size
+        else:
+            # Load dimensions from pkl file
+            dims = self.asset_dimensions[asset_type]
+            size = [float(dims[0]), float(dims[1]), float(dims[2])]  # [width, depth, height]
+            original_asset_size = size.copy()  # Store the original asset dimensions
+        
+        cuboid = {
+            'id': obj_id,
+            'description': description,
+            'type': asset_type,  # Store the asset type
+            'position': [0.0, 0.0, 0.0],  # Place on ground (z = height/2)
+            'size': size,
+            'original_asset_size': original_asset_size,  # Store original asset dimensions
+            'azimuth': 0.0,
+            'color': self._get_next_color()
+        }
+        self.objects.append(cuboid)
+        return obj_id
+
+
+    def update_cuboid(self, obj_id, x, y, z, azimuth, width, depth, height):
+        if 0 <= obj_id < len(self.objects):
+            obj = self.objects[obj_id]
+            obj['position'] = [x, y, z]
+            obj['size'] = [width, depth, height]
+            # Note: We do NOT update original_asset_size here - it stays unchanged
+            obj['azimuth'] = azimuth
+            return True
+        return False
+
+
+    def delete_cuboid(self, obj_id):
+        if 0 <= obj_id < len(self.objects):
+            del self.objects[obj_id]
+            # Update IDs for remaining objects
+            for i, obj in enumerate(self.objects):
+                obj['id'] = i
+            return True
+        return False
+            
+    def set_camera_elevation(self, elevation_deg):
+        assert type(elevation_deg) == float or type(elevation_deg) == int, f"{type(elevation_deg) = }" 
+        self.camera_elevation = np.clip(elevation_deg, 0.0, 90.0)
+        return f"Camera elevation set to {elevation_deg}°"
+
+    def set_camera_lens(self, lens_value):
+        self.camera_lens = np.clip(lens_value, 10.0, 200.0)
+        return f"Camera lens set to {lens_value}mm"
+
+    def set_surrounding_prompt(self, prompt):  # Add this method
+        self.surrounding_prompt = prompt
+        return f"Surrounding prompt updated"
+
+    def _convert_to_blender_format(self):
+        """Convert internal objects format to server expected format"""
+        subjects_data = []
+        
+        for obj in self.objects:
+            subject_data = {
+                'subject_name': obj['description'],
+                'x': float(obj['position'][0]),
+                'y': float(obj['position'][1]), 
+                'z': float(obj['position'][2]),
+                'azimuth': float(obj['azimuth']),
+                'width': float(obj['size'][0]),
+                'depth': float(obj['size'][1]),
+                'height': float(obj['size'][2]),
+                'base_color': obj['color']
+            }
+            subjects_data.append(subject_data)
+            
+        camera_data = {
+            'camera_elevation': float(np.radians(self.camera_elevation)),
+            'lens': float(self.camera_lens),
+            'global_scale': 1.0
+        }
+        
+        return subjects_data, camera_data
+
+    def render_cv_view(self, subjects_data: list, camera_data: dict) -> Image.Image:
+        """Render only the CV view."""
+        if not subjects_data:
+            return Image.new('RGB', (512, 512), color='gray')
+        
+        return self.render_client._send_render_request(self.render_client.cv_server_url, subjects_data, camera_data)
+
+
+    def render_scene(self, width=512, height=512):
+        """Render only CV view using the render client."""
+        print(f"calling render_scene")
+        if not self.objects:
+            # Return empty image if no objects
+            empty_cv = Image.new('RGB', (width, height), color='gray')
+            return empty_cv
+        
+        # Convert to server expected format
+        subjects_data, camera_data = self._convert_to_blender_format()
+        print(f"passing {subjects_data = } to render_cv_view in SceneManager") 
+        
+        # Render CV view only
+        cv_img = self.render_cv_view(subjects_data, camera_data)
+        
+        return cv_img
+    
+# --- Gradio Interface Logic ---
+scene_manager = SceneManager()
+
+def get_cuboid_list_html():
+    """Generate HTML for the cuboid list with position-based colors"""
+    if not scene_manager.objects:
+        return "<div style='text-align: center; padding: 20px; color: #888;'>No cuboids yet. Add one to get started!</div>"
+    
+    html = "<div style='display: flex; flex-direction: column; gap: 8px;'>"
+    for obj_idx, obj in enumerate(scene_manager.objects):
+        # Get position-based color
+        # x, y = obj['position'][0], obj['position'][1]
+        # rgb_color = map_point_to_rgb(x, y)
+        rgb_color = COLORS[obj_idx % len(COLORS)]  
+        hex_color = rgb_to_hex(rgb_color)
+        
+        # Create a lighter version for gradient end
+        lighter_rgb = tuple(min(1.0, c + 0.2) for c in rgb_color)
+        lighter_hex = rgb_to_hex(lighter_rgb)
+        
+        html += f"""
+        <div style='background: linear-gradient(135deg, {hex_color} 0%, {lighter_hex} 100%); 
+                    padding: 12px; border-radius: 8px; color: white; text-shadow: 1px 1px 2px rgba(0,0,0,0.5);'>
+            <div style='font-weight: bold; font-size: 14px;'>{obj['description']}</div>
+            <div style='font-size: 11px; opacity: 0.9; margin-top: 4px;'>
+                Pos: ({obj['position'][0]:.1f}, {obj['position'][1]:.1f}, {obj['position'][2]:.1f}) | 
+                Size: {obj['size'][0]:.1f}×{obj['size'][1]:.1f}×{obj['size'][2]:.1f}
+            </div>
+        </div>
+        """
+    html += "</div>"
+    return html
+
+
+def add_cuboid_event(description_input, asset_type, camera_elevation, camera_lens):
+    """Add a new cuboid"""
+    if not description_input.strip():
+        description_input = "New Cuboid"
+    
+    new_id = scene_manager.add_cuboid(description_input, asset_type)
+    cv_img = scene_manager.render_scene()
+    
+    # Create choices for radio buttons
+    choices = [f"{obj['description']}" for obj in scene_manager.objects]
+    
+    # Get the new object data
+    new_obj = scene_manager.objects[new_id]
+    
+    return (
+        gr.update(value=""),  # Clear description input
+        gr.update(value="Custom"),  # Reset type dropdown to Custom
+        cv_img,
+        get_cuboid_list_html(),
+        gr.update(choices=choices, value=new_obj['description']),  # Radio with new selection
+        gr.update(visible=True),  # Show editor
+        gr.update(value=new_obj['description']),  # Set description in editor
+        gr.update(value=new_obj['position'][0]),
+        gr.update(value=new_obj['position'][1]),
+        gr.update(value=new_obj['position'][2]),
+        gr.update(value=new_obj['azimuth']),
+        gr.update(value=new_obj['size'][0]),
+        gr.update(value=new_obj['size'][1]),
+        gr.update(value=new_obj['size'][2]),
+        gr.update(value=1.0)  # Reset scale to 1.0
+    )
+
+
+def select_cuboid_event(selected_name):
+    """When a cuboid is selected from radio buttons"""
+    if not selected_name:
+        return [gr.update(visible=False)] + [gr.update() for _ in range(9)]  # Changed from 8 to 9
+    
+    # Find the cuboid by description
+    obj = None
+    for o in scene_manager.objects:
+        if o['description'] == selected_name:
+            obj = o
+            break
+    
+    if obj is None:
+        return [gr.update(visible=False)] + [gr.update() for _ in range(9)]
+    
+    return (
+        gr.update(visible=True),  # Show editor
+        gr.update(value=obj['description']),  # Set description
+        gr.update(value=obj['position'][0]),
+        gr.update(value=obj['position'][1]),
+        gr.update(value=obj['position'][2]),
+        gr.update(value=obj['azimuth']),
+        gr.update(value=obj['size'][0]),
+        gr.update(value=obj['size'][1]),
+        gr.update(value=obj['size'][2]),
+        gr.update(value=1.0)  # Reset scale to 1.0
+    )
+
+
+def delete_selected_cuboid(selected_name, camera_elevation, camera_lens):
+    """Delete the currently selected cuboid"""
+    if not selected_name:
+        return gr.update(), get_cuboid_list_html(), gr.update(), gr.update(visible=False)
+    
+    # Find and delete the cuboid
+    obj_id = None
+    for i, obj in enumerate(scene_manager.objects):
+        if obj['description'] == selected_name:
+            obj_id = i
+            break
+    
+    if obj_id is not None:
+        scene_manager.delete_cuboid(obj_id)
+    
+    cv_img = scene_manager.render_scene()
+    
+    # Update choices
+    choices = [f"{obj['description']}" for obj in scene_manager.objects]
+    
+    return (
+        cv_img,
+        get_cuboid_list_html(),
+        gr.update(choices=choices, value=None),
+        gr.update(visible=False)
+    )
+
+
+def update_cuboid_event(selected_name, camera_elevation, camera_lens, description, x, y, z, azimuth, width, depth, height, scale):
+    """Update the selected cuboid including description and scale"""
+    scene_manager.set_camera_elevation(camera_elevation)
+    scene_manager.set_camera_lens(camera_lens)
+    
+    if selected_name:
+        # Find the cuboid by description
+        obj_id = None
+        for i, obj in enumerate(scene_manager.objects):
+            if obj['description'] == selected_name:
+                obj_id = i
+                break
+        
+        if obj_id is not None:
+            # Update description first if changed
+            if description.strip() and description.strip() != selected_name:
+                scene_manager.update_cuboid_description(obj_id, description.strip())
+            
+            # Apply scale to dimensions
+            scaled_width = width * scale
+            scaled_depth = depth * scale
+            scaled_height = height * scale
+            
+            # Update other properties with scaled dimensions
+            scene_manager.update_cuboid(obj_id, x, y, z, azimuth, scaled_width, scaled_depth, scaled_height)
+            
+            # Get updated object for return
+            updated_obj = scene_manager.objects[obj_id]
+            new_name = updated_obj['description']
+    
+    cv_img = scene_manager.render_scene()
+    
+    # Update choices with new descriptions
+    choices = [f"{obj['description']}" for obj in scene_manager.objects]
+    
+    # Return updated HTML, image, radio choices, new selection, updated sliders, and reset scale
+    return (
+        get_cuboid_list_html(), 
+        cv_img,
+        gr.update(choices=choices, value=new_name if obj_id is not None else None),
+        gr.update(value=scaled_width if obj_id is not None else width),   # Update width slider
+        gr.update(value=scaled_depth if obj_id is not None else depth),   # Update depth slider
+        gr.update(value=scaled_height if obj_id is not None else height), # Update height slider
+        gr.update(value=1.0)  # Reset scale to 1.0
+    )
+
+
+def camera_change_event(camera_elevation, camera_lens):
+    """Handle camera control changes"""
+    scene_manager.set_camera_elevation(camera_elevation)
+    scene_manager.set_camera_lens(camera_lens)
+    cv_img = scene_manager.render_scene()
+    return cv_img
+
+
+def surrounding_prompt_change_event(prompt_text):  # Add this function
+    """Handle surrounding prompt changes"""
+    scene_manager.set_surrounding_prompt(prompt_text)
+    return None  # No visual update needed
+
+
+def render_segmask_event(camera_elevation, camera_lens, surrounding_prompt):
+    """Render segmentation masks for all objects"""
+    if not scene_manager.objects:
+        return "⚠️ No objects to render", gr.update(visible=False), []
+    
+    # Get subject descriptions
+    subject_descriptions = [obj['description'] for obj in scene_manager.objects]
+    
+    # Now you have access to:
+    # - surrounding_prompt: the text from surrounding_prompt_input
+    # - subject_descriptions: list of all subject descriptions
+    
+    print(f"Surrounding prompt: {surrounding_prompt}")
+    print(f"Subject descriptions: {subject_descriptions}")
+
+    placeholder_prompt = "a photo of PLACEHOLDER " + surrounding_prompt 
+
+    # Create placeholder text
+    subject_embeds = [] 
+    for subject_idx, subject_desc in enumerate(subject_descriptions): 
+        input_ids = tokenizer.encode(subject_desc, return_tensors="pt", max_length=77)[0] 
+        subject_embed = {"input_ids_t5": input_ids.tolist()} 
+        subject_embeds.append(subject_embed)
+
+    placeholder_text = ""
+    for subject in subject_descriptions[:-1]:
+        placeholder_text = placeholder_text + f"<placeholder> {subject} and "
+    for subject in subject_descriptions[-1:]:
+        placeholder_text = placeholder_text + f"<placeholder> {subject}"
+    placeholder_text = placeholder_text.strip()
+
+    placeholder_token_prompt = placeholder_prompt.replace("PLACEHOLDER", placeholder_text) 
+
+    call_ids = get_call_ids_from_placeholder_prompt_flux(prompt=placeholder_token_prompt, 
+                                             subjects=subject_descriptions, 
+                                             subjects_embeds=subject_embeds,
+                                             debug=True
+                                            ) 
+    print(f"Generated call IDs: {call_ids}") 
+
+    
+    # Convert to server expected format
+    subjects_data, camera_data = scene_manager._convert_to_blender_format()
+    
+    # You can add the prompt and descriptions to the request if needed
+    # For example, add to subjects_data or camera_data before sending
+    
+    # Render segmentation masks
+    success, segmask_images, error_msg = scene_manager.render_client.render_segmasks(subjects_data, camera_data)
+
+    # copy all the data to the correct location  
+    root_save_dir = "/archive/vaibhav.agrawal/a-bev-of-the-latents/gradio_files/" 
+    os.system("rm /archive/vaibhav.agrawal/a-bev-of-the-latents/gradio_files/*") 
+    shutil.move("cv_render.jpg", osp.join(root_save_dir, "cv_render.jpg")) 
+    for subject_idx in range(len(subject_descriptions)): 
+        shutil.move(f"{str(subject_idx).zfill(3)}_segmask_cv.png", osp.join(root_save_dir, f"main__segmask_{str(subject_idx).zfill(3)}__{1.00}.png")) 
+    
+    jsonl = [{
+        "cv": osp.join(root_save_dir, "cv_render.jpg"),
+        "target": osp.join(root_save_dir, "cv_render.jpg"),
+        "cuboids_segmasks": [osp.join(root_save_dir, f"main__segmask_{str(subject_idx).zfill(3)}__{1.00}.png") for subject_idx in range(len(subject_descriptions))],
+        "PLACEHOLDER_prompts": placeholder_prompt, 
+        "subjects": subject_descriptions, 
+        "call_ids": call_ids, 
+    }] 
+
+    with open(osp.join(root_save_dir, "cuboids.jsonl"), "w") as f: 
+        for item in jsonl: 
+            f.write(json.dumps(item) + "\n") 
+    
+    if success:
+        return (
+            f"✅ Successfully rendered {len(segmask_images)} segmentation masks",
+            gr.update(visible=True),
+            segmask_images
+        )
+    else:
+        return (
+            f"❌ Failed to render segmentation masks: {error_msg}",
+            gr.update(visible=False),
+            []
+        )
+
+
+def harmonize_event(selected_name, camera_elevation, camera_lens):
+    """Harmonize all object scales and update the scene"""
+    message = scene_manager.harmonize_scales()
+    print(message)
+    
+    cv_img = scene_manager.render_scene()
+    
+    # If a cuboid is selected, update its sliders
+    if selected_name:
+        obj = None
+        for o in scene_manager.objects:
+            if o['description'] == selected_name:
+                obj = o
+                break
+        
+        if obj is not None:
+            return (
+                cv_img,
+                get_cuboid_list_html(),
+                gr.update(value=obj['position'][0]),
+                gr.update(value=obj['position'][1]),
+                gr.update(value=obj['position'][2]),
+                gr.update(value=obj['azimuth']),
+                gr.update(value=obj['size'][0]),
+                gr.update(value=obj['size'][1]),
+                gr.update(value=obj['size'][2])
+            )
+    
+    # No object selected or object not found
+    return (
+        cv_img,
+        get_cuboid_list_html(),
+        gr.update(),
+        gr.update(),
+        gr.update(),
+        gr.update(),
+        gr.update(),
+        gr.update(),
+        gr.update()
+    )
+
+
+def save_scene_event():
+    """Save the current scene to a pkl file"""
+    success, filepath, error = scene_manager.save_scene_to_pkl()
+    
+    if success:
+        return f"✅ Scene saved successfully to: {filepath}\n📋 Saved parameters: {scene_manager.inference_params}"
+    else:
+        return f"❌ Failed to save scene: {error}"
+
+
+def load_scene_event(filepath):
+    """Load a scene from a pkl file and restore all parameters"""
+    if not filepath.strip():
+        return (
+            "⚠️ Please enter a file path",
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),  # surrounding_prompt
+            gr.update(),  # checkpoint
+            gr.update(),  # height
+            gr.update(),  # width
+            gr.update(),  # seed
+            gr.update(),  # guidance
+            gr.update()   # steps
+        )
+    
+    success, num_objects, error = scene_manager.load_scene_from_pkl(filepath)
+    
+    if success:
+        # Re-render the scene
+        cv_img = scene_manager.render_scene()
+        
+        # Update UI components
+        choices = [f"{obj['description']}" for obj in scene_manager.objects]
+        
+        params_msg = f"✅ Scene loaded: {num_objects} objects\n📋 Restored parameters: {scene_manager.inference_params}"
+        
+        return (
+            params_msg,
+            cv_img,
+            get_cuboid_list_html(),
+            gr.update(choices=choices, value=None),
+            gr.update(visible=False),
+            gr.update(value=scene_manager.camera_elevation),
+            gr.update(value=scene_manager.camera_lens),
+            gr.update(value=scene_manager.surrounding_prompt),
+            gr.update(value=scene_manager.inference_params['checkpoint']),
+            gr.update(value=scene_manager.inference_params['height']),
+            gr.update(value=scene_manager.inference_params['width']),
+            gr.update(value=scene_manager.inference_params['seed']),
+            gr.update(value=scene_manager.inference_params['guidance_scale']),
+            gr.update(value=scene_manager.inference_params['num_inference_steps'])
+        )
+    else:
+        return (
+            f"❌ {error}",
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update()
+        )
+
+
+# --- Gradio UI Layout ---
+with gr.Blocks(
+    theme=gr.themes.Soft(
+        primary_hue="green",
+        secondary_hue="gray",
+        neutral_hue="gray"
+    ),
+    css="""
+    .gradio-container {
+        background: linear-gradient(135deg, #0d1117 0%, #1a3d2e 50%, #000000 100%) !important;
+        color: #ffffff !important;
+    }
+    .block {
+        background: rgba(15, 36, 25, 0.8) !important;
+        border: 1px solid #2d5a41 !important;
+        border-radius: 8px !important;
+    }
+    .form {
+        background: rgba(15, 36, 25, 0.6) !important;
+    }
+    h1, h2, h3, h4, h5, h6 {
+        color: #ffffff !important;
+    }
+    .markdown {
+        color: #e6e6e6 !important;
+    }
+    label {
+        color: #cccccc !important;
+    }
+    .gr-button {
+        background: linear-gradient(135deg, #2d5a41, #3d6a51) !important;
+        border: 1px solid #4a7c59 !important;
+        color: #ffffff !important;
+    }
+    .gr-button:hover {
+        background: linear-gradient(135deg, #3d6a51, #4a7c59) !important;
+    }
+    .gr-input, .gr-textbox, .gr-dropdown {
+        background: rgba(15, 36, 25, 0.8) !important;
+        border: 1px solid #2d5a41 !important;
+        color: #ffffff !important;
+    }
+    .gr-input:focus, .gr-textbox:focus {
+        border-color: #4a7c59 !important;
+        background: rgba(26, 61, 46, 0.8) !important;
+    }
+    .gr-slider input[type="range"] {
+        background: #2d5a41 !important;
+    }
+    .gr-slider input[type="range"]::-webkit-slider-thumb {
+        background: #4a7c59 !important;
+    }
+    .gr-radio label {
+        color: #cccccc !important;
+    }
+    .gr-panel {
+        background: rgba(15, 36, 25, 0.6) !important;
+        border: 1px solid #2d5a41 !important;
+    }
+    """
+) as demo:
+    gr.Markdown("# [CVPR-2026] 3D Aware Occlusion Control in Text-to-Image Generation 🏞️🧱")
+    # TOP ROW
+    with gr.Row():
+        # TOP LEFT - Edit Properties
+        with gr.Column(scale=1):
+            # Add description textbox at the top
+            # with gr.Column(visible=False) as editor_section:
+            #     gr.Markdown("## ✏️ Edit Properties")
+                
+            #     delete_btn = gr.Button("❌ Delete Selected Cuboid", variant="stop", size="sm")
+                
+            #     with gr.Row():
+            #         edit_x = gr.Slider(-10, 10, value=0, step=0.1, label="X")
+            #         edit_y = gr.Slider(-10, 10, value=0, step=0.1, label="Y")
+            #         edit_z = gr.Slider(0, 10, value=1, step=0.1, label="Z")
+                
+            #     edit_azimuth = gr.Slider(-180, 180, value=0, step=1, label="Azimuth (°)")
+                
+            #     with gr.Row():
+            #         edit_width = gr.Slider(0.1, 5, value=1, step=0.1, label="Width")
+            #         edit_depth = gr.Slider(0.1, 5, value=1, step=0.1, label="Depth")
+            #         edit_height = gr.Slider(0.1, 5, value=1, step=0.1, label="Height")
+            with gr.Column(visible=False) as editor_section:
+                gr.Markdown("## ✏️ Edit Properties")
+
+                edit_description = gr.Textbox(
+                    label="Description",
+                    placeholder="Enter object description",
+                    info="Description cannot be empty"
+                )
+                
+                delete_btn = gr.Button("❌ Delete Selected Cuboid", variant="stop", size="sm")
+                
+                with gr.Row():
+                    edit_x = gr.Slider(-10, 10, value=0, step=0.1, label="X")
+                    edit_y = gr.Slider(-10, 10, value=0, step=0.1, label="Y")
+                    edit_z = gr.Slider(0, 10, value=1, step=0.1, label="Z")
+                
+                edit_azimuth = gr.Slider(-180, 180, value=0, step=1, label="Azimuth (°)")
+                
+                with gr.Row():
+                    edit_width = gr.Slider(0.1, 5, value=1, step=0.1, label="Width")
+                    edit_depth = gr.Slider(0.1, 5, value=1, step=0.1, label="Depth")
+                    edit_height = gr.Slider(0.1, 5, value=1, step=0.1, label="Height")
+                
+                # Add scale slider
+                edit_scale = gr.Slider(
+                    0.1, 3.0, value=1.0, step=0.1, 
+                    label="Scale",
+                    info="Multiplier for all dimensions (resets to 1.0 after update)"
+                )
+                
+                # Add the Update Scene button
+                update_scene_btn = gr.Button("🔄 Update Scene", variant="primary", size="sm")
+
+        # TOP MIDDLE - Camera View
+        with gr.Column(scale=1):
+            gr.Markdown("## 📷 Camera View")
+            cv_image_output = gr.Image(label="Camera View", height=400)
+
+        # TOP RIGHT - Generated Image
+        with gr.Column(scale=1):
+            gr.Markdown("## 🎨 Generated Image")
+            generated_image_output = gr.Image(label="Generated Image", height=400)
+
+    # BOTTOM ROW
+    with gr.Row():
+        # BOTTOM LEFT - Cuboid List and Selection
+        with gr.Column(scale=1):
+            gr.Markdown("## 📦 Cuboids")
+            cuboid_list_html = gr.HTML(get_cuboid_list_html())
+            
+            gr.Markdown("### Select Cuboid to Edit")
+            cuboid_radio = gr.Radio(choices=[], label="", visible=True)
+
+        # BOTTOM RIGHT - Camera Controls and Add New Cuboid
+        with gr.Column(scale=2):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Global Controls")
+                    camera_elevation_slider = gr.Slider(0, 90, value=30, label="Camera Elevation (degrees)")
+                    camera_lens_slider = gr.Slider(10, 200, value=50, label="Camera Lens (mm)")
+
+                    # Add surrounding prompt textbox
+                    surrounding_prompt_input = gr.Textbox(
+                        placeholder="e.g., in a forest, in a city, on a beach",
+                        label="Surrounding Prompt",
+                        info="Describe the surrounding environment"
+                    )
+                    
+                    gr.Markdown("## 🔧 Scene Tools")
+                    harmonize_btn = gr.Button("⚖️ Harmonize Scales", variant="secondary")
+                    
+                    # Save/Load Section
+                    gr.Markdown("## 💾 Save/Load Scene")
+                    with gr.Row():
+                        save_scene_btn = gr.Button("💾 Save Scene", variant="secondary")
+                        load_scene_btn = gr.Button("📂 Load Scene", variant="secondary")
+                    
+                    load_path_input = gr.Textbox(
+                        placeholder="/path/to/scene.pkl",
+                        label="Load Scene Path",
+                        info="Enter path to pkl file to load"
+                    )
+                    save_load_status = gr.Markdown("")
+                
+                with gr.Column():
+                    gr.Markdown("## ➕ Add New Cuboid")
+                    add_cuboid_description_input = gr.Textbox(placeholder="Enter cuboid description", label="Description")
+                    asset_type_dropdown = gr.Dropdown(
+                        choices=scene_manager.get_asset_type_choices(),
+                        value="Custom",
+                        label="Type",
+                        info="Select asset type to load dimensions, or choose Custom"
+                    )
+                    add_cuboid_btn = gr.Button("Add Cuboid", variant="primary")
+                    generate_btn = gr.Button("🎨 Generate Image", variant="primary")
+                    
+                    # Add checkpoint dropdown
+                    checkpoint_dropdown = gr.Dropdown(
+                        choices=CHECKPOINT_NAMES,
+                        value=CHECKPOINT_NAMES[0] if CHECKPOINT_NAMES else None,
+                        label="Checkpoint",
+                        info="Select model checkpoint for generation"
+                    )
+
+                    # Inference Parameters
+                    gr.Markdown("### Inference Parameters")
+                    
+                    with gr.Row():
+                        inference_height = gr.Slider(
+                            minimum=256, maximum=1024, value=512, step=64,
+                            label="Height"
+                        )
+                        inference_width = gr.Slider(
+                            minimum=256, maximum=1024, value=512, step=64,
+                            label="Width"
+                        )
+                    
+                    inference_seed = gr.Number(
+                        value=42, label="Random Seed", precision=0
+                    )
+                    
+                    inference_guidance = gr.Slider(
+                        minimum=1.0, maximum=10.0, value=3.5, step=0.5,
+                        label="Guidance Scale"
+                    )
+                    
+                    inference_steps = gr.Slider(
+                        minimum=10, maximum=50, value=25, step=1,
+                        label="Inference Steps"
+                    )
+    
+    # Event Handlers
+    def add_cuboid_with_auto_update(description_input, asset_type, camera_elevation, camera_lens):
+        """Add cuboid and auto-update scene"""
+        result = add_cuboid_event(description_input, asset_type, camera_elevation, camera_lens)
+        return result
+    
+    # Update add_cuboid_btn.click event handler (around line 850):
+    add_cuboid_btn.click(
+        add_cuboid_with_auto_update,
+        inputs=[add_cuboid_description_input, asset_type_dropdown, camera_elevation_slider, camera_lens_slider],
+        outputs=[
+            add_cuboid_description_input,
+            asset_type_dropdown,
+            cv_image_output,
+            cuboid_list_html,
+            cuboid_radio,
+            editor_section,
+            edit_description,
+            edit_x, edit_y, edit_z,
+            edit_azimuth,
+            edit_width, edit_depth, edit_height,
+            edit_scale  # Add this
+        ]
+    )
+    
+    # Update the cuboid_radio.change event handler (around line 860):
+    cuboid_radio.change(
+        select_cuboid_event,
+        inputs=[cuboid_radio],
+        outputs=[
+            editor_section,
+            edit_description,
+            edit_x, edit_y, edit_z,
+            edit_azimuth,
+            edit_width, edit_depth, edit_height,
+            edit_scale  # Add this
+        ]
+    )
+    
+    delete_btn.click(
+        delete_selected_cuboid,
+        inputs=[cuboid_radio, camera_elevation_slider, camera_lens_slider],
+        outputs=[cv_image_output, cuboid_list_html, cuboid_radio, editor_section]
+    )
+
+    # Save/Load handlers
+    save_scene_btn.click(
+        save_scene_event,
+        inputs=[],
+        outputs=[save_load_status]
+    )
+    
+    load_scene_btn.click(
+        load_scene_event,
+        inputs=[load_path_input],
+        outputs=[
+            save_load_status,
+            cv_image_output,
+            cuboid_list_html,
+            cuboid_radio,
+            editor_section,
+            camera_elevation_slider,
+            camera_lens_slider,
+            surrounding_prompt_input,
+            checkpoint_dropdown,
+            inference_height,
+            inference_width,
+            inference_seed,
+            inference_guidance,
+            inference_steps
+        ]
+    )
+    
+    # Auto-update scene when sliders change
+    # for slider in [edit_x, edit_y, edit_z, edit_azimuth, edit_width, edit_depth, edit_height]:
+    #     slider.change(
+    #         update_cuboid_event,
+    #         inputs=[
+    #             cuboid_radio,
+    #             camera_elevation_slider,
+    #             camera_lens_slider,
+    #             edit_x, edit_y, edit_z,
+    #             edit_azimuth,
+    #             edit_width, edit_depth, edit_height
+    #         ],
+    #         outputs=[cuboid_list_html, cv_image_output]
+    #     )
+    # Update the update_scene_btn.click event handler (around line 920):
+    update_scene_btn.click(
+        update_cuboid_event,
+        inputs=[
+            cuboid_radio,
+            camera_elevation_slider,
+            camera_lens_slider,
+            edit_description,
+            edit_x, edit_y, edit_z,
+            edit_azimuth,
+            edit_width, edit_depth, edit_height,
+            edit_scale  # Add this
+        ],
+        outputs=[
+            cuboid_list_html, 
+            cv_image_output, 
+            cuboid_radio,
+            edit_width,   # Add this
+            edit_depth,   # Add this
+            edit_height,  # Add this
+            edit_scale    # Add this (to reset to 1.0)
+        ]
+    )
+
+
+    # Update generate button click handler
+    generate_btn.click(
+        generate_image_event,
+        inputs=[
+            camera_elevation_slider, 
+            camera_lens_slider, 
+            surrounding_prompt_input, 
+            checkpoint_dropdown,
+            inference_height,
+            inference_width,
+            inference_seed,
+            inference_guidance,
+            inference_steps
+        ],
+        outputs=[save_load_status, cv_image_output, generated_image_output]
+    )
+
+    
+    harmonize_btn.click(
+        harmonize_event,
+        inputs=[cuboid_radio, camera_elevation_slider, camera_lens_slider],
+        outputs=[
+            cv_image_output,
+            cuboid_list_html,
+            edit_x, edit_y, edit_z,
+            edit_azimuth,
+            edit_width, edit_depth, edit_height
+        ]
+    )
+    
+    # Camera controls
+    for control in [camera_elevation_slider, camera_lens_slider]:
+        control.change(
+            camera_change_event,
+            inputs=[camera_elevation_slider, camera_lens_slider],
+            outputs=[cv_image_output]
+        )
+
+    # Surrounding prompt control
+    surrounding_prompt_input.change(
+        surrounding_prompt_change_event,
+        inputs=[surrounding_prompt_input],
+        outputs=[]
+    )
+
+
+    # Initial render
+    def initial_render():
+        cv_img = scene_manager.render_scene()
+        gen_img = Image.new('RGB', (512, 512), color='white')
+        return cv_img, gen_img
+    
+    demo.load(
+        initial_render, 
+        outputs=[cv_image_output, generated_image_output]
+    )
+
+
+if __name__ == "__main__":
+    import os 
+    os.system("./launch_blender_backend.sh &")
+    # Initialize inference engine (load model once at startup)
+    initialize_inference_engine(base_model_path="black-forest-labs/FLUX.1-dev")
+    demo.launch(share=True)

gradio_app/blender_backend.py

@@ -0,0 +1,1521 @@
+import bpy 
+import bpy_extras 
+import numpy as np 
+import bmesh
+import copy 
+import PIL 
+from PIL import Image 
+import matplotlib.pyplot as plt 
+import colorsys
+import os 
+import os.path as osp 
+import shutil 
+import sys 
+import math 
+import mathutils 
+import random 
+import cv2 
+from object_scales import scales 
+import matplotlib.colors as mcolors
+import torch 
+
+def map_point_to_rgb(x, y, z):
+    """
+    Map (x, y) inside the frustum to an RGB color with continuity and variation.
+    """
+    # Frustum boundaries
+    X_MIN, X_MAX = -12.0, -1.0
+    Y_MIN_AT_XMIN, Y_MAX_AT_XMIN = -4.5, 4.5
+    Y_MIN_AT_XMAX, Y_MAX_AT_XMAX = -0.5, 0.5
+    Z_MIN, Z_MAX = 0.0, 2.50   
+    # Normalize x to [0, 1]
+    x_norm = (x - X_MIN) / (X_MAX - X_MIN)
+    x_norm = np.clip(x_norm, 0, 1)
+
+    # Compute current Y bounds at given x using linear interpolation
+    y_min = Y_MIN_AT_XMIN + x_norm * (Y_MIN_AT_XMAX - Y_MIN_AT_XMIN)
+    y_max = Y_MAX_AT_XMIN + x_norm * (Y_MAX_AT_XMAX - Y_MAX_AT_XMIN)
+
+    # Normalize y to [0, 1] within current bounds
+    if y_max != y_min:
+        y_norm = (y - y_min) / (y_max - y_min)
+    else:
+        y_norm = 0.5
+    y_norm = np.clip(y_norm, 0, 1)
+
+    z_norm = (z - Z_MIN) / (Z_MAX - Z_MIN) 
+
+    # Color mapping: more variation along x
+    r = x_norm
+    # g = 0.5 * y_norm + 0.25 * x_norm
+    g = y_norm 
+    b = z_norm 
+
+    return (r, g, b)
+
+
+def set_world_color(color=(0.1, 0.1, 0.1)):
+    """
+    Sets the world background color to match the grid floor.
+    
+    Args:
+        color (tuple): RGB color values (0-1 range)
+    """
+    scene = bpy.context.scene
+    
+    # Create a new world if it doesn't exist
+    if scene.world is None:
+        world = bpy.data.worlds.new(name="World")
+        scene.world = world
+    else:
+        world = scene.world
+    
+    # Enable use of nodes for the world
+    world.use_nodes = True
+    
+    # Get the node tree
+    nodes = world.node_tree.nodes
+    links = world.node_tree.links
+    
+    # Find or create the Background node
+    background_node = None
+    for node in nodes:
+        if node.type == 'BACKGROUND':
+            background_node = node
+            break
+    
+    if background_node is None:
+        # Clear existing nodes and create new ones
+        nodes.clear()
+        background_node = nodes.new(type='ShaderNodeBackground')
+        output_node = nodes.new(type='ShaderNodeOutputWorld')
+        links.new(background_node.outputs['Background'], output_node.inputs['Surface'])
+    
+    # Set the background color
+    background_node.inputs['Color'].default_value = (*color, 1.0)
+    background_node.inputs['Strength'].default_value = 1.0
+
+
+COLORS = [
+    (1.0, 0.0, 0.0),    # Red
+    (0.0, 0.8, 0.2),    # Green
+    (0.0, 0.0, 1.0),    # Blue
+    (1.0, 1.0, 0.0),    # Yellow
+    (0.0, 1.0, 1.0),    # Cyan
+    (1.0, 0.0, 1.0),    # Magenta
+    (1.0, 0.6, 0.0),    # Orange
+    (0.6, 0.0, 0.8),    # Purple
+    (0.0, 0.4, 0.0),    # Dark Green
+    (0.8, 0.8, 0.8),    # Light Gray
+    (0.2, 0.2, 0.2)     # Dark Gray
+]
+
+def do_z_pass(seg_masks: torch.Tensor, dist_values: torch.Tensor) -> torch.Tensor:
+    """
+    Performs a z-pass on segmentation masks based on distance values to the camera.
+    For each pixel, if multiple subjects' masks are active, only the one with the smallest distance (closest) remains active.
+    
+    Args:
+        seg_masks (torch.Tensor): Binary segmentation masks of shape (n_subjects, h, w) with dtype uint8.
+        dist_values (torch.Tensor): Distance values for each subject of shape (n_subjects,).
+    
+    Returns:
+        torch.Tensor: Processed segmentation masks after z-pass, same shape and dtype as seg_masks.
+    """
+    # Ensure tensors are on the same device
+    device = seg_masks.device
+    
+    # Get dimensions
+    n_subjects, h, w = seg_masks.shape
+    
+    # Reshape distance values for broadcasting across spatial dimensions
+    dist_values_expanded = dist_values.view(n_subjects, 1, 1)
+    
+    # Create a tensor where active pixels have their distance, others have a high value (1e10)
+    masked_dist = torch.where(seg_masks.bool(), dist_values_expanded, torch.tensor(1e10, device=device))
+    
+    # Find the subject index with the minimum distance for each pixel (shape (h, w))
+    closest_indices = torch.argmin(masked_dist, dim=0)
+    
+    # Initialize output tensor with zeros
+    output = torch.zeros_like(seg_masks)
+    
+    # Scatter 1s into the output tensor where the closest subject's indices are
+    # closest_indices.unsqueeze(0) adds a dummy dimension to match scatter's expected shape
+    output.scatter_(
+        dim=0,
+        index=closest_indices.unsqueeze(0),
+        src=torch.ones_like(closest_indices.unsqueeze(0), dtype=output.dtype)
+    )
+    
+    # Zero out any positions where the original mask was inactive
+    output = output * seg_masks
+    
+    return output
+
+
+def get_image_to_world_matrix(camera_obj, render):
+    """
+    Calculates the matrix to transform a point from clip space to world space.
+
+    Args:
+        camera_obj (bpy.types.Object): The camera object.
+        render (bpy.types.RenderSettings): The scene's render settings.
+
+    Returns:
+        mathutils.Matrix: The 4x4 matrix for clip-to-world transformation.
+    """
+    # Get the camera's view matrix (world to camera)
+    view_matrix = camera_obj.matrix_world.inverted()
+
+    # Get the camera's projection matrix
+    # This matrix depends on the render resolution, so it's best to calculate it
+    # for the specific dimensions you're using.
+    projection_matrix = camera_obj.calc_matrix_camera(
+        bpy.context.evaluated_depsgraph_get(),
+        x=render.resolution_x,
+        y=render.resolution_y,
+        scale_x=render.pixel_aspect_x,
+        scale_y=render.pixel_aspect_y,
+    )
+
+    # Combine and invert to get the clip-to-world matrix
+    clip_to_world_matrix = (projection_matrix @ view_matrix).inverted()
+    
+    return clip_to_world_matrix
+
+
+def unproject_image_point(camera_obj, image_coord, depth):
+    """
+    Transforms a 2D image coordinate with a depth value into a 3D world coordinate.
+
+    Args:
+        camera_obj (bpy.types.Object): The camera used for rendering.
+        image_coord (tuple or list): The (x, y) pixel coordinate.
+        depth (float): The depth value at that coordinate (from the Z-pass).
+
+    Returns:
+        mathutils.Vector: The calculated 3D point in world space.
+    """
+    render = bpy.context.scene.render
+    
+    # 1. Get the clip-to-world transformation matrix
+    clip_to_world_mat = get_image_to_world_matrix(camera_obj, render)
+
+    # 2. Convert image coordinates to Normalized Device Coordinates (NDC)
+    #    (from [0, res] to [-1, 1])
+    ndc_x = (image_coord[0] / render.resolution_x) * 2 - 1
+    ndc_y = (image_coord[1] / render.resolution_y) * 2 - 1
+
+    # In Blender's Z-pass, the depth value is the distance from the camera's plane.
+    # We can use Blender's utility function to find the 3D vector for the pixel.
+    # This vector is in camera space and points from the camera towards the pixel.
+    view_vector = bpy_extras.view3d_utils.region_2d_to_vector_3d(
+        bpy.context.region,
+        bpy.context.space_data.region_3d,
+        image_coord
+    )
+
+    # 4. Project the view vector into world space and scale by depth
+    # The view_vector is normalized and in camera space.
+    # To get the point in world space, we transform the vector by the camera's
+    # world matrix (not the view matrix).
+    world_vector = camera_obj.matrix_world.to_3x3() @ view_vector
+    
+    # The depth from the Z-pass is the distance along the camera's local Z-axis.
+    # To find the true distance along the ray, we must account for the angle.
+    # We can calculate the scaling factor 't' for our world_vector.
+    camera_forward = -camera_obj.matrix_world.col[2].xyz
+    t = depth / world_vector.dot(camera_forward)
+
+    # 5. Calculate the final world coordinate
+    # Start from the camera's location and move along the ray.
+    world_point = camera_obj.matrix_world.translation + (t * world_vector)
+
+    return world_point
+
+# --- Example Usage ---
+# This example assumes you have an active scene with a camera and have rendered an image.
+# You would typically run this after rendering, where you can access the depth map.
+
+
+def multiply_random_color(obj, random_color):
+    """
+    Multiplies the existing base color of an object's materials
+    with a random color.
+    """
+    for material_slot in obj.material_slots:
+        if material_slot.material:
+            material = material_slot.material
+            if material.use_nodes:
+                nodes = material.node_tree.nodes
+                links = material.node_tree.links
+
+                # Find the Principled BSDF node
+                principled_bsdf = nodes.get("Principled BSDF")
+                if not principled_bsdf:
+                    continue
+
+                # Get the node connected to the Base Color input
+                base_color_input = principled_bsdf.inputs.get("Base Color")
+                if not base_color_input:
+                    continue
+
+                # Create a MixRGB node and set it to multiply
+                mix_rgb_node = nodes.new(type='ShaderNodeMixRGB')
+                mix_rgb_node.blend_type = 'MULTIPLY'
+                mix_rgb_node.inputs['Fac'].default_value = 2.00  
+                mix_rgb_node.location = (principled_bsdf.location.x - 200, principled_bsdf.location.y)
+
+                # Set the second color to a random color
+                mix_rgb_node.inputs['Color2'].default_value = random_color
+
+                # If a node is already connected to the Base Color,
+                # connect it to the first color input of the MixRGB node.
+                if base_color_input.is_linked:
+                    original_link = base_color_input.links[0]
+                    original_node = original_link.from_node
+                    original_socket = original_link.from_socket
+                    links.new(original_node.outputs[original_socket.name], mix_rgb_node.inputs['Color1'])
+                    links.remove(original_link)
+                else:
+                    # If no node is connected, use the original default color
+                    original_color = base_color_input.default_value
+                    mix_rgb_node.inputs['Color1'].default_value = original_color
+
+                # Connect the MixRGB node to the Principled BSDF's Base Color
+                links.new(mix_rgb_node.outputs['Color'], base_color_input)
+
+
+OUTPUT_DIR = "four_subject_renders" 
+OBJECTS_DIR = "obja_2units_along_y/glbs" 
+
+NUM_AZIMUTH_BINS = 1      
+NUM_LIGHTS = 1    
+
+MAX_TRIES = 25   
+
+IMG_DIM = 1024 
+
+MASK_RES = 50  
+
+THRESHOLD_LOWER = 150 
+THRESHOLD_UPPER = 768  
+
+ROOT_OBJS_DIR = "/ssd_scratch/vaibhav.agrawal/a-bev-of-the-latents/glb_files/"  
+
+OBJ_SIDE_LENGTH = 2.0 
+
+def calculate_iou(box1, box2):
+    """
+    Calculate the Intersection over Union (IoU) of two bounding boxes.
+
+    Parameters:
+    box1, box2: Each box is defined by a tuple (x1, y1, x2, y2)
+                where (x1, y1) is the top-left corner and (x2, y2) is the bottom-right corner.
+
+    Returns:
+    float: IoU value
+    """
+    # Unpack coordinatesO
+    x1_min, y1_min, x1_max, y1_max = box1
+    x2_min, y2_min, x2_max, y2_max = box2
+    
+    # Determine the coordinates of the intersection rectangle
+    inter_x_min = max(x1_min, x2_min)
+    inter_y_min = max(y1_min, y2_min)
+    inter_x_max = min(x1_max, x2_max)
+    inter_y_max = min(y1_max, y2_max)
+    
+    # Compute the area of intersection rectangle
+    inter_width = max(0, inter_x_max - inter_x_min)
+    inter_height = max(0, inter_y_max - inter_y_min)
+    intersection_area = inter_width * inter_height
+    
+    # Compute the area of both bounding boxes
+    box1_area = (x1_max - x1_min) * (y1_max - y1_min)
+    box2_area = (x2_max - x2_min) * (y2_max - y2_min)
+    
+    # Compute the area of the union
+    union_area = box1_area + box2_area - intersection_area
+    
+    # Compute IoU
+    iou = intersection_area / union_area if union_area > 0 else 0
+    
+    return iou
+
+
+def get_object_2d_bbox(empty_obj, scene):
+    """
+    Get the 2D bounding box coordinates of an object in the rendered image.
+    
+    Args:
+        empty_obj (bpy.types.Object): The empty object containing the child mesh objects.
+        scene (bpy.types.Scene): The current scene.
+        
+    Returns:
+        tuple: A tuple containing the 2D bounding box coordinates in pixel space
+              in the format (min_x, min_y, max_x, max_y).
+    """
+    # Get the render settings
+    render = scene.render
+    res_x = render.resolution_x
+    res_y = render.resolution_y
+    
+    # Initialize the bounding box coordinates
+    min_x, min_y = float('inf'), float('inf')
+    max_x, max_y = float('-inf'), float('-inf')
+
+    depsgraph = bpy.context.evaluated_depsgraph_get()
+
+    # Iterate through the child mesh objects
+    for obj in empty_obj.children:
+        if obj.type == 'MESH':
+            # Get the bounding box coordinates in world space
+            bbox_corners = [obj.matrix_world @ mathutils.Vector(corner) for corner in obj.bound_box]
+            
+            # Transform the bounding box corners to camera space
+            for corner in bbox_corners:
+                corner_2d = bpy_extras.object_utils.world_to_camera_view(scene, scene.camera, corner)
+
+                # Scale the coordinates to pixel space
+                x = corner_2d.x * res_x
+                y = (1 - corner_2d.y) * res_y  # Flip Y since Blender renders from bottom to top
+                
+                # Update the bounding box coordinates
+                min_x = min(min_x, x)
+                min_y = min(min_y, y)
+                max_x = max(max_x, x)
+                max_y = max(max_y, y)
+    
+    # Return the 2D bounding box coordinates in pixel space
+    return (int(min_x), int(min_y), int(max_x), int(max_y))
+
+def reset_cameras(scene) -> None:
+    """Resets the cameras in the scene to a single default camera."""
+    # Delete all existing cameras
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="CAMERA")
+    bpy.ops.object.delete()
+    
+    # Create a new camera with default properties
+    bpy.ops.object.camera_add()
+    
+    # Get the camera by searching for it (it will be the only camera)
+    new_camera = None
+    for obj in scene.objects:
+        if obj.type == 'CAMERA':
+            new_camera = obj
+            break
+    
+    new_camera.name = "Camera"
+    
+    # Set the new camera as the active camera for the scene
+    scene.camera = new_camera
+
+
+def add_plane():
+    print(f"in add_plane")
+    
+    # Create mesh data
+    mesh = bpy.data.meshes.new("Plane")
+    backdrop = bpy.data.objects.new("Plane", mesh)
+    bpy.context.scene.collection.objects.link(backdrop)
+    
+    # Create plane geometry using bmesh
+    bm = bmesh.new()
+    bmesh.ops.create_grid(bm, x_segments=1, y_segments=1, size=25.0)  # size=25 gives 50x50 plane
+    bm.to_mesh(mesh)
+    bm.free()
+    
+    # Add material
+    mat_backdrop = bpy.data.materials.new(name="WhiteMaterial")
+    mat_backdrop.diffuse_color = (0, 0, 0, 1)  # Black
+    backdrop.data.materials.append(mat_backdrop)
+
+
+def add_plane_cycles():
+    print(f"in add_plane")
+    
+    # Create mesh data
+    mesh = bpy.data.meshes.new("Plane")
+    backdrop = bpy.data.objects.new("Plane", mesh)
+    bpy.context.scene.collection.objects.link(backdrop)
+    
+    # Create plane geometry using bmesh
+    bm = bmesh.new()
+    bmesh.ops.create_grid(bm, x_segments=1, y_segments=1, size=25.0)  # size=25 gives 50x50 plane
+    bm.to_mesh(mesh)
+    bm.free()
+    
+    # Add material
+    mat_backdrop = bpy.data.materials.new(name="WhiteMaterial")
+    mat_backdrop.diffuse_color = (0.050, 0.050, 0.050, 1)  # White
+    backdrop.data.materials.append(mat_backdrop)
+
+
+def remove_all_planes():
+    # Deselect all objects first
+    bpy.ops.object.select_all(action='DESELECT')
+
+    # Select all plane objects in the scene
+    for obj in bpy.data.objects:
+        if obj.type == 'MESH' and obj.name.startswith('Plane'):
+            obj.select_set(True)
+    
+    # Delete all selected planes
+    bpy.ops.object.delete()
+
+
+def remove_all_lights():
+    """Remove all lights from the scene without using operators."""
+    lights_to_remove = [obj for obj in bpy.data.objects if obj.type == 'LIGHT']
+    
+    for light in lights_to_remove:
+        bpy.data.objects.remove(light, do_unlink=True)
+    
+    # Clean up orphaned light data blocks
+    for light_data in bpy.data.lights:
+        if light_data.users == 0:
+            bpy.data.lights.remove(light_data)
+
+
+def set_lights_cv(radius, center, num_points, intensity):
+    print(f"in set_lights_cv") 
+    radius = radius + 10.0 
+    phi = np.random.uniform(-np.pi / 2, np.pi / 2, num_points)         # azimuthal angle
+    cos_theta = np.random.uniform(0.50, 1.0, num_points)          # cos of polar angle
+    theta = np.arccos(cos_theta)                              # polar angle
+    x = np.sin(theta) * np.cos(phi)
+    y = np.sin(theta) * np.sin(phi)
+    z = cos_theta  # cos(theta) == z on unit sphere
+    # Scale to radius and shift to center
+    points = np.stack([x, y, z], axis=1) * radius + center
+    for point in points: 
+        # Track objects before adding light
+        before_objs = set(bpy.data.objects)
+        bpy.ops.object.light_add(type='POINT', location=point)  
+        after_objs = set(bpy.data.objects)
+        
+        # Get the newly created light
+        diff_objs = after_objs - before_objs
+        light = list(diff_objs)[0]
+        
+        light.data.energy = intensity
+        light.data.use_shadow = True   
+        # light.data.shadow_soft_size = 1.0  # Adjust shadow softness if needed 
+    return points
+
+
+def adjust_color_brightness(rgb_color, factor):
+    """
+    Adjusts the brightness of an RGB color by a multiplicative factor.
+
+    Args:
+        rgb_color (tuple): The base color as an (R, G, B) or (R, G, B, A) tuple.
+        factor (float): The factor to multiply the brightness by.
+                        > 1.0 makes it lighter, < 1.0 makes it darker.
+
+    Returns:
+        tuple: The new (R, G, B, A) color.
+    """
+    # Use only RGB for conversion, keep alpha separate
+    h, s, v = colorsys.rgb_to_hsv(rgb_color[0], rgb_color[1], rgb_color[2])
+    
+    # Multiply the Value (brightness) by the factor, and clamp it between 0 and 1
+    v = max(0, min(1, v * factor))
+    
+    new_rgb = colorsys.hsv_to_rgb(h, s, v)
+    
+    # Return as an RGBA tuple, preserving original alpha if it exists
+    alpha = rgb_color[3] if len(rgb_color) == 4 else 1.0
+    return (new_rgb[0], new_rgb[1], new_rgb[2], alpha)
+
+
+def get_primitive_object_translucent(base_color=(0.0, 1.0, 0.0), edge_color=None, face_opacity=0.025):
+    """
+    Spawns a cuboid primitive with individually colored faces and highlighted edges.
+
+    Args:
+        base_color (tuple): The base RGB color for the faces.
+        edge_color (tuple): The RGBA color for the edges (defaults to white).
+        face_opacity (float): The opacity of the cuboid faces (0.0 = invisible, 1.0 = opaque). Default is 0.2.
+    """
+    # --- Create the Cuboid and Parent ---
+    bpy.ops.object.empty_add(type="PLAIN_AXES")
+    # empty_object = bpy.context.object
+    empty_object = bpy.data.objects.new("Empty", None)
+    before_objs = set(bpy.data.objects)
+    bpy.ops.mesh.primitive_cube_add(size=0.5, location=(0, 0, 0))
+    after_objs = set(bpy.data.objects)
+    diff_objs = after_objs - before_objs
+
+    obj = None
+    for o in diff_objs:
+        obj = o
+        obj.parent = empty_object
+        world_matrix = obj.matrix_world
+        obj.matrix_world = world_matrix
+
+    # --- Create and Assign Materials for Each Face ---
+    if obj:
+        # left front right back bottom top 
+        brightness_factors = [
+            0.30, 0.30, 0.30, 0.30, 1.00, 0.30, 
+        ]
+        colors = [adjust_color_brightness(base_color, factor) for factor in brightness_factors]
+
+        for i, color in enumerate(colors):
+            material = bpy.data.materials.new(name=f"FaceColor_{i}")
+            material.use_nodes = True
+            obj.data.materials.append(material)
+
+            nodes = material.node_tree.nodes
+            links = material.node_tree.links
+            nodes.clear()
+
+            # Create Principled BSDF instead of Emission for proper transparency
+            bsdf = nodes.new(type="ShaderNodeBsdfPrincipled")
+            bsdf.location = (0, 0)
+            bsdf.inputs['Base Color'].default_value = color
+            bsdf.inputs['Alpha'].default_value = face_opacity  # Set face opacity
+            bsdf.inputs['Emission Color'].default_value = color[:3] + (1.0,)  # Fixed: Use 'Emission Color' instead of 'Emission'
+            bsdf.inputs['Emission Strength'].default_value = 1.0  # Emission strength
+            
+            material_output = nodes.new(type="ShaderNodeOutputMaterial")
+            material_output.location = (200, 0)
+            links.new(bsdf.outputs['BSDF'], material_output.inputs['Surface'])
+            
+            # Enable transparency settings for the material
+            material.blend_method = 'BLEND'
+            material.show_transparent_back = False
+
+        if len(obj.data.polygons) == len(colors):
+            for i, poly in enumerate(obj.data.polygons):
+                poly.material_index = i
+        else:
+            print("Warning: The number of colors does not match the number of faces.")
+
+        # --- Add Wireframe Edges ---
+        # edge_material = bpy.data.materials.new(name="EdgeDelimiterMaterial")
+        # edge_material.use_nodes = True
+        
+        # nodes = edge_material.node_tree.nodes
+        # links = edge_material.node_tree.links
+        # nodes.clear()
+
+        # if edge_color is None: 
+        #     edge_color = adjust_color_brightness(base_color, 0.10) 
+
+        # edge_emission_node = nodes.new(type="ShaderNodeEmission")
+        # edge_emission_node.inputs['Color'].default_value = edge_color
+        # edge_output_node = nodes.new(type="ShaderNodeOutputMaterial")
+        # links.new(edge_emission_node.outputs['Emission'], edge_output_node.inputs['Surface'])
+
+        # obj.data.materials.append(edge_material)
+        
+        # wire_mod = obj.modifiers.new(name="EdgeDelimiter", type='WIREFRAME')
+        # wire_mod.thickness = 0.01
+        # wire_mod.use_replace = False
+        # wire_mod.material_offset = len(obj.data.materials) - 1
+
+    # --- Bounding Box Calculation ---
+    bbox_corners = []
+    bpy.context.view_layer.update()
+    for child in empty_object.children:
+        for corner in child.bound_box:
+            world_corner = child.matrix_world @ mathutils.Vector(corner)
+            bbox_corners.append(world_corner)
+
+    if not bbox_corners:
+        return 0, empty_object
+
+    min_x = min(corner.x for corner in bbox_corners)
+    min_y = min(corner.y for corner in bbox_corners)
+    min_z = min(corner.z for corner in bbox_corners)
+
+    max_x = max(corner.x for corner in bbox_corners)
+    max_y = max(corner.y for corner in bbox_corners)
+    max_z = max(corner.z for corner in bbox_corners)
+
+    return max_z, empty_object
+
+
+def get_primitive_object_translucent_rgb(base_color=(0.0, 1.0, 0.0), edge_color=None, face_opacity=0.025):
+    """
+    Spawns a cuboid primitive with individually colored faces and highlighted edges.
+
+    Args:
+        base_color (tuple): The base RGB color for the faces.
+        edge_color (tuple): The RGBA color for the edges (defaults to white).
+        face_opacity (float): The opacity of the cuboid faces (0.0 = invisible, 1.0 = opaque). Default is 0.2.
+    """
+    # --- Create the Cuboid and Parent ---
+    bpy.ops.object.empty_add(type="PLAIN_AXES")
+    # empty_object = bpy.context.object
+    empty_object = bpy.data.objects.new("Empty", None)
+    before_objs = set(bpy.data.objects)
+    bpy.ops.mesh.primitive_cube_add(size=0.5, location=(0, 0, 0))
+    after_objs = set(bpy.data.objects)
+    diff_objs = after_objs - before_objs
+
+    obj = None
+    for o in diff_objs:
+        obj = o
+        obj.parent = empty_object
+        world_matrix = obj.matrix_world
+        obj.matrix_world = world_matrix
+
+    # --- Create and Assign Materials for Each Face ---
+    if obj:
+        # left front right back bottom top 
+        brightness_factors = [
+            0.50, 0.50, 0.50, 0.50, 0.50, 0.50, 
+        ]
+        red = (1.0, 0.0, 0.0, 1.0) 
+        green = (0.0, 1.0, 0.0, 1.0)  
+        blue = (0.0, 0.0, 1.0, 1.0) 
+        colors = [adjust_color_brightness(green, factor) for factor in brightness_factors[:4]] + [adjust_color_brightness(blue, brightness_factors[4])] + [adjust_color_brightness(red, brightness_factors[5])] 
+        colors = [colors[-2], colors[-1], colors[0], colors[1], colors[2], colors[3]] 
+
+        for i, color in enumerate(colors):
+            material = bpy.data.materials.new(name=f"FaceColor_{i}")
+            material.use_nodes = True
+            obj.data.materials.append(material)
+
+            nodes = material.node_tree.nodes
+            links = material.node_tree.links
+            nodes.clear()
+
+            # Create Principled BSDF instead of Emission for proper transparency
+            bsdf = nodes.new(type="ShaderNodeBsdfPrincipled")
+            bsdf.location = (0, 0)
+            bsdf.inputs['Base Color'].default_value = color
+            bsdf.inputs['Alpha'].default_value = face_opacity  # Set face opacity
+            bsdf.inputs['Emission Color'].default_value = color[:3] + (1.0,)  # Fixed: Use 'Emission Color' instead of 'Emission'
+            bsdf.inputs['Emission Strength'].default_value = 1.0  # Emission strength
+            
+            material_output = nodes.new(type="ShaderNodeOutputMaterial")
+            material_output.location = (200, 0)
+            links.new(bsdf.outputs['BSDF'], material_output.inputs['Surface'])
+            
+            # Enable transparency settings for the material
+            material.blend_method = 'BLEND'
+            material.show_transparent_back = False
+
+        if len(obj.data.polygons) == len(colors):
+            for i, poly in enumerate(obj.data.polygons):
+                poly.material_index = i
+        else:
+            print("Warning: The number of colors does not match the number of faces.")
+
+        # --- Add Wireframe Edges ---
+        edge_material = bpy.data.materials.new(name="EdgeDelimiterMaterial")
+        edge_material.use_nodes = True
+        
+        nodes = edge_material.node_tree.nodes
+        links = edge_material.node_tree.links
+        nodes.clear()
+
+        if edge_color is None: 
+            edge_color = adjust_color_brightness(base_color, 0.10) 
+
+        edge_emission_node = nodes.new(type="ShaderNodeEmission")
+        edge_emission_node.inputs['Color'].default_value = edge_color
+        edge_output_node = nodes.new(type="ShaderNodeOutputMaterial")
+        links.new(edge_emission_node.outputs['Emission'], edge_output_node.inputs['Surface'])
+
+        obj.data.materials.append(edge_material)
+        
+        wire_mod = obj.modifiers.new(name="EdgeDelimiter", type='WIREFRAME')
+        wire_mod.thickness = 0.01
+        wire_mod.use_replace = False
+        wire_mod.material_offset = len(obj.data.materials) - 1
+
+    # --- Bounding Box Calculation ---
+    bbox_corners = []
+    bpy.context.view_layer.update()
+    for child in empty_object.children:
+        for corner in child.bound_box:
+            world_corner = child.matrix_world @ mathutils.Vector(corner)
+            bbox_corners.append(world_corner)
+
+    if not bbox_corners:
+        return 0, empty_object
+
+    min_x = min(corner.x for corner in bbox_corners)
+    min_y = min(corner.y for corner in bbox_corners)
+    min_z = min(corner.z for corner in bbox_corners)
+
+    max_x = max(corner.x for corner in bbox_corners)
+    max_y = max(corner.y for corner in bbox_corners)
+    max_z = max(corner.z for corner in bbox_corners)
+
+    return max_z, empty_object
+
+
+
+def get_primitive_object(base_color=(0.0, 1.0, 0.0), edge_color=None):
+    """
+    Spawns a cuboid primitive with individually colored faces and highlighted edges.
+
+    Args:
+        base_color (tuple): The base RGB color for the faces.
+        edge_color (tuple): The RGBA color for the edges (defaults to white).
+    """
+    # --- Create the Empty Parent ---
+    empty_object = bpy.data.objects.new("Empty", None)
+    bpy.context.scene.collection.objects.link(empty_object)
+    empty_object.empty_display_type = 'PLAIN_AXES'
+    
+    # --- Create the Cuboid using bmesh ---
+    mesh = bpy.data.meshes.new("Cube")
+    obj = bpy.data.objects.new("Cube", mesh)
+    bpy.context.scene.collection.objects.link(obj)
+    
+    # Create cube geometry
+    bm = bmesh.new()
+    bmesh.ops.create_cube(bm, size=0.5)
+    bm.to_mesh(mesh)
+    bm.free()
+    
+    # Set parent
+    obj.parent = empty_object
+    world_matrix = obj.matrix_world
+    obj.matrix_world = world_matrix
+
+    # --- Create and Assign Materials for Each Face ---
+    if obj:
+        # left front right back bottom top 
+        brightness_factors = [
+            0.35, 0.20, 0.65, 0.90, 0.50, 0.50   
+        ]
+        colors = [adjust_color_brightness(base_color, factor) for factor in brightness_factors]
+
+        for i, color in enumerate(colors):
+            material = bpy.data.materials.new(name=f"FaceColor_{i}")
+            material.use_nodes = True
+            obj.data.materials.append(material)
+
+            nodes = material.node_tree.nodes
+            links = material.node_tree.links
+            nodes.clear()
+
+            emission_node = nodes.new(type="ShaderNodeEmission")
+            emission_node.inputs['Color'].default_value = color
+            material_output = nodes.new(type="ShaderNodeOutputMaterial")
+            links.new(emission_node.outputs['Emission'], material_output.inputs['Surface'])
+            
+            material.blend_method = 'BLEND'
+            material.show_transparent_back = False
+
+        if len(obj.data.polygons) == len(colors):
+            for i, poly in enumerate(obj.data.polygons):
+                poly.material_index = i
+        else:
+            print("Warning: The number of colors does not match the number of faces.")
+
+        # --- MODIFICATION START: Add White Edges ---
+
+        # 1. Create a new material for the wireframe edges
+        edge_material = bpy.data.materials.new(name="EdgeDelimiterMaterial")
+        edge_material.use_nodes = True
+        
+        # Set up the nodes for a simple white emission shader
+        nodes = edge_material.node_tree.nodes
+        links = edge_material.node_tree.links
+        nodes.clear()
+
+        if edge_color is None: 
+            edge_color = adjust_color_brightness(base_color, 0.10) 
+
+        edge_emission_node = nodes.new(type="ShaderNodeEmission")
+        edge_emission_node.inputs['Color'].default_value = edge_color
+        edge_output_node = nodes.new(type="ShaderNodeOutputMaterial")
+        links.new(edge_emission_node.outputs['Emission'], edge_output_node.inputs['Surface'])
+
+        # 2. Add the edge material to the object's material slots
+        obj.data.materials.append(edge_material)
+        
+        # 3. Add and configure the Wireframe modifier
+        wire_mod = obj.modifiers.new(name="EdgeDelimiter", type='WIREFRAME')
+        wire_mod.thickness = 0.01  # The thickness of the edge lines
+        wire_mod.use_replace = False  # Set to False to keep the original faces
+        # This offset tells the modifier to use the last material we added (the white one)
+        wire_mod.material_offset = len(obj.data.materials) - 1
+
+        # --- MODIFICATION END ---
+
+
+    # --- Bounding Box Calculation (remains the same) ---
+    bbox_corners = []
+    # Update the dependency graph to ensure modifiers are accounted for
+    bpy.context.view_layer.update()
+    for child in empty_object.children:
+        # Use child.bound_box which is in object's local space
+        for corner in child.bound_box:
+            # Convert corner to world space
+            world_corner = child.matrix_world @ mathutils.Vector(corner)
+            bbox_corners.append(world_corner)
+
+    if not bbox_corners:
+        return 0, empty_object # Return a default value if no corners found
+
+    min_x = min(corner.x for corner in bbox_corners)
+    min_y = min(corner.y for corner in bbox_corners)
+    min_z = min(corner.z for corner in bbox_corners)
+
+    max_x = max(corner.x for corner in bbox_corners)
+    max_y = max(corner.y for corner in bbox_corners)
+    max_z = max(corner.z for corner in bbox_corners)
+
+    return max_z, empty_object
+
+class BlenderCuboidRenderer:
+    def __init__(self, render_engine): 
+        """
+        Initialize the Blender cuboid renderer.
+        
+        Args:
+            img_dim (int): Image dimensions (square)
+            render_engine (str): Blender render engine ('EEVEE' or 'CYCLES')
+            num_lights (int): Number of lights to add
+            max_tries (int): Maximum tries for placement
+        """
+        self.img_dim = 1024 
+        self.render_engine = render_engine
+        self.blender_grid_dims = scales
+
+        self.radius = 6.0 
+        self.center = -6.0 
+        
+        # Scene references
+        self.context = None
+        self.scene = None
+        self.camera = None
+        self.render = None
+
+        # Setup the scene
+        self.setup_scene()
+
+        
+    def setup_scene(self):
+        """
+        Setup the basic Blender scene with camera, lighting, and render settings.
+        
+        Args:
+            camera_data (dict): Camera configuration containing elevation, lens, global_scale, etc.
+        """
+        # Get all objects in the scene
+        objects_to_remove = []
+        
+        for obj in bpy.data.objects:
+            # Remove default cube, plane, camera, and lights
+            if obj.type in {'MESH', 'LIGHT', 'CAMERA'}:
+                objects_to_remove.append(obj)
+        
+        # Delete the objects
+        for obj in objects_to_remove:
+            bpy.data.objects.remove(obj, do_unlink=True)
+        
+        # Also clear orphaned data
+        for mesh in bpy.data.meshes:
+            if mesh.users == 0:
+                bpy.data.meshes.remove(mesh)
+        
+        for light in bpy.data.lights:
+            if light.users == 0:
+                bpy.data.lights.remove(light)
+        
+        for camera in bpy.data.cameras:
+            if camera.users == 0:
+                bpy.data.cameras.remove(camera)
+
+        bpy.context.scene.world = None 
+
+        # Initialize Blender scene
+        # bpy.ops.wm.read_factory_settings(use_empty=True)
+        self.context = bpy.context 
+        self.scene = self.context.scene 
+        if self.render_engine == "CYCLES":  
+            self.scene.cycles.samples = 32 
+        self.render = self.scene.render 
+        
+        # Set render engine and resolution
+        self.render.engine = self.render_engine  
+        self.context.scene.render.resolution_x = self.img_dim  
+        self.context.scene.render.resolution_y = self.img_dim  
+        self.context.scene.render.resolution_percentage = 100  
+
+        # Setup compositing nodes
+        self._setup_compositing()
+        
+        
+    def _setup_compositing(self):
+        """Setup Blender compositing nodes for depth and RGB output."""
+        self.context.scene.use_nodes = True
+        tree = self.context.scene.node_tree
+        links = tree.links
+
+        self.context.scene.render.use_compositing = True 
+        self.context.view_layer.use_pass_z = True
+        
+        # clear default nodes
+        for n in tree.nodes:
+            tree.nodes.remove(n)
+        
+        # create input render layer node
+        rl = tree.nodes.new('CompositorNodeRLayers')      
+
+        map_node = tree.nodes.new(type="CompositorNodeMapValue")
+        map_node.size = [0.05]
+        map_node.use_min = True
+        map_node.min = [0]
+        map_node.use_max = True
+        map_node.max = [65336]
+        links.new(rl.outputs[2], map_node.inputs[0])
+
+        invert = tree.nodes.new(type="CompositorNodeInvert")
+        links.new(map_node.outputs[0], invert.inputs[1])
+        
+        # create output node
+        v = tree.nodes.new('CompositorNodeViewer')   
+        v.use_alpha = True 
+
+        # create a file output node and set the path
+        fileOutput = tree.nodes.new(type="CompositorNodeOutputFile")
+        fileOutput.base_path = "."
+        links.new(invert.outputs[0], fileOutput.inputs[0])
+
+        # Links
+        links.new(rl.outputs[0], v.inputs[0])  # link Image to Viewer Image RGB
+        links.new(rl.outputs['Depth'], v.inputs[1])  # link Render Z to Viewer Image Alpha
+
+        # Update scene to apply changes
+        self.context.view_layer.update() 
+
+
+    def _setup_camera_cv(self, camera_data):
+        """Setup camera position and orientation."""
+        reset_cameras(self.scene) 
+        self.camera = self.scene.objects["Camera"] 
+        
+        elevation = camera_data["camera_elevation"] 
+        tan_elevation = np.tan(elevation) 
+        cos_elevation = np.cos(elevation) 
+        sin_elevation = np.sin(elevation) 
+
+        radius = self.radius 
+        center = self.center 
+        
+        self.camera.location = mathutils.Vector((radius * cos_elevation + center, 0, radius * sin_elevation))  
+        direction = mathutils.Vector((-1, 0, -tan_elevation)) 
+        self.context.scene.camera = self.camera 
+        rot_quat = direction.to_track_quat("-Z", "Y") 
+        self.camera.rotation_euler = rot_quat.to_euler() 
+        self.camera.data.lens = camera_data["lens"]  
+
+    def _create_cuboid_objects_translucent(self, subjects_data, opacity=0.025):
+        """Create primitive cuboid objects for all subjects."""
+        for subject_idx, subject_data in enumerate(subjects_data): 
+            # rgb_color = map_point_to_rgb(x, y) 
+            rgb_color = COLORS[subject_idx % len(COLORS)] 
+            _, prim_obj = get_primitive_object_translucent(base_color=rgb_color, face_opacity=opacity) 
+            prim_obj.location = np.array([100, 0, 0]) 
+            subject_data["prim_obj"] = prim_obj
+
+    def _create_cuboid_objects_translucent_rgb(self, subjects_data, opacity=0.025):
+        """Create primitive cuboid objects for all subjects."""
+        for subject_idx, subject_data in enumerate(subjects_data): 
+            x = subject_data["x"][0] 
+            y = subject_data["y"][0] 
+            z = subject_data["z"][0] 
+            base_color = map_point_to_rgb(x, y, z)
+            _, prim_obj = get_primitive_object_translucent_rgb(base_color=base_color, face_opacity=opacity) 
+            prim_obj.location = np.array([100, 0, 0]) 
+            subject_data["prim_obj"] = prim_obj
+
+            
+    def _place_objects(self, subjects_data, camera_data):
+        """Place objects in the scene according to their data."""
+        global_scale = camera_data["global_scale"] 
+        
+        for subject_data in subjects_data: 
+            x = subject_data["x"][0]  
+            y = subject_data["y"][0]  
+            z = global_scale * subject_data["dims"][2] / 2.0 + subject_data["z"][0]  
+            subject_data["prim_obj"].location = np.array([x, y, z])  
+            subject_data["prim_obj"].scale = global_scale * np.array(subject_data["dims"]) * 2.0 
+            subject_data["prim_obj"].rotation_euler[2] = subject_data["azimuth"][0]  
+
+    def render_cv(self, subjects_data, camera_data, num_samples=1, output_path="main.jpg"):
+        """
+        Main render method that takes subjects data and renders the scene.
+        
+        Args:
+            subjects_data (list): List of subject dictionaries containing position, dims, etc.
+            camera_data (dict): Camera configuration
+            num_samples (int): Number of samples to render (currently only supports 1)
+            output_path (str): Path to save the rendered image
+            
+        Returns:
+            None
+        """
+        center = (-6.0, 0.0, 0.0) 
+        radius = 6.0 
+
+        print(f"render_cv received {subjects_data = }") 
+
+        # print(f"render_cv received {subjects_data = }")
+        for subject_data in subjects_data: 
+            subject_data["azimuth"][0] = np.deg2rad(subject_data["azimuth"][0]) 
+            subject_data["x"][0] = subject_data["x"][0] + center[0] 
+            subject_data["y"][0] = subject_data["y"][0] + center[1] 
+            subject_data["z"][0] = subject_data["z"][0] + center[2] 
+        # Setup camera
+        self._setup_camera_cv(camera_data)
+        
+        set_lights_cv(self.radius, np.array([self.center, 0, 0]), 20, intensity=7000.0)
+        
+        # Add ground plane
+        add_plane()
+
+        assert num_samples == 1, "for now, only implemented for a single sample"  
+        assert "global_scale" in camera_data.keys(), "global_scale must be set for EEVEE" 
+        
+        # Create primitive objects for subjects
+        self._create_cuboid_objects_translucent(subjects_data, opacity=0.025)
+        # self._create_cuboid_objects(subjects_data)
+        
+        # Place objects in scene
+        self._place_objects(subjects_data, camera_data)
+        
+        # Perform rendering
+        print(f"SUCCESS, rendering...")
+        self.context.scene.render.filepath = output_path 
+        self.context.scene.render.image_settings.file_format = "JPEG" 
+        bpy.ops.render.render(write_still=True) 
+        
+        print(f"Rendered scene saved to: {output_path}")
+
+        self.cleanup() 
+
+    def render_final_representation(self, subjects_data, camera_data, num_samples=1, output_path="main.jpg"):
+        """
+        Main render method that takes subjects data and renders the scene.
+        
+        Args:
+            subjects_data (list): List of subject dictionaries containing position, dims, etc.
+            camera_data (dict): Camera configuration
+            num_samples (int): Number of samples to render (currently only supports 1)
+            output_path (str): Path to save the rendered image
+            
+        Returns:
+            None
+        """
+        assert self.render.engine == "CYCLES", "render_final_representation only works with CYCLES render engine" 
+        center = (-6.0, 0.0, 0.0) 
+        radius = 6.0 
+
+        print(f"render_cv received {subjects_data = }") 
+
+        # print(f"render_cv received {subjects_data = }")
+        for subject_data in subjects_data: 
+            subject_data["azimuth"][0] = np.deg2rad(subject_data["azimuth"][0]) 
+            subject_data["x"][0] = subject_data["x"][0] + center[0] 
+            subject_data["y"][0] = subject_data["y"][0] + center[1] 
+            subject_data["z"][0] = subject_data["z"][0] + center[2] 
+        # Setup camera
+        self._setup_camera_cv(camera_data)
+        
+        print(f"setting lights in cycles...")
+        set_lights_cv(self.radius, np.array([self.center, 0, 0]), 5, intensity=700.0)
+        
+        # Add ground plane
+        print(f"adding plane in cycles...")
+        add_plane_cycles()
+
+        assert num_samples == 1, "for now, only implemented for a single sample"  
+        assert "global_scale" in camera_data.keys(), "global_scale must be set for EEVEE" 
+        
+        # Create primitive objects for subjects
+        self._create_cuboid_objects_translucent_rgb(subjects_data, opacity=0.025)
+        # self._create_cuboid_objects(subjects_data)
+        
+        # Place objects in scene
+        self._place_objects(subjects_data, camera_data)
+        
+        # Perform rendering
+        print(f"SUCCESS, rendering...")
+        self.context.scene.render.filepath = output_path 
+        self.context.scene.render.image_settings.file_format = "JPEG" 
+        bpy.ops.render.render(write_still=True) 
+        
+        print(f"Rendered scene saved to: {output_path}")
+
+        self.cleanup() 
+
+
+    def render_paper_figure(self, subjects_data, camera_data, num_samples=1, output_path="main.jpg"):
+        """
+        Main render method that takes subjects data and renders the scene.
+        
+        Args:
+            subjects_data (list): List of subject dictionaries containing position, dims, etc.
+            camera_data (dict): Camera configuration
+            num_samples (int): Number of samples to render (currently only supports 1)
+            output_path (str): Path to save the rendered image
+            
+        Returns:
+            None
+        """
+        assert self.render.engine == "CYCLES", "render_final_representation only works with CYCLES render engine" 
+        center = (-6.0, 0.0, 0.0) 
+        radius = 6.0 
+
+        print(f"render_cv received {subjects_data = }") 
+
+        set_world_color((1.0, 1.0, 1.0))  # white background
+
+        # print(f"render_cv received {subjects_data = }")
+        for subject_data in subjects_data: 
+            subject_data["azimuth"][0] = np.deg2rad(subject_data["azimuth"][0]) 
+            subject_data["x"][0] = subject_data["x"][0] + center[0] 
+            subject_data["y"][0] = subject_data["y"][0] + center[1] 
+            subject_data["z"][0] = subject_data["z"][0] + center[2] 
+        # Setup camera
+        self._setup_camera_cv(camera_data)
+        
+        print(f"setting lights in cycles...")
+        set_lights_cv(self.radius, np.array([self.center, 0, 0]), 5, intensity=7000.0)
+        
+        # Add ground plane
+        print(f"adding plane in cycles...")
+
+        assert num_samples == 1, "for now, only implemented for a single sample"  
+        assert "global_scale" in camera_data.keys(), "global_scale must be set for EEVEE" 
+        
+        # Create primitive objects for subjects
+        self._create_cuboid_objects_translucent(subjects_data, opacity=0.35)
+        # self._create_cuboid_objects(subjects_data)
+        
+        # Place objects in scene
+        self._place_objects(subjects_data, camera_data)
+        
+        # Perform rendering
+        print(f"SUCCESS, rendering...")
+        self.context.scene.render.filepath = output_path 
+        self.context.scene.render.image_settings.file_format = "JPEG" 
+        bpy.ops.render.render(write_still=True) 
+        
+        print(f"Rendered scene saved to: {output_path}")
+
+        self.cleanup() 
+
+
+    def cleanup(self):
+        """Clean up the scene for next render."""
+        # Remove all lights
+        remove_all_lights()
+        
+        # Remove all other objects (meshes, empties, etc.)
+        objects_to_remove = [obj for obj in bpy.data.objects]
+        
+        for obj in objects_to_remove:
+            bpy.data.objects.remove(obj, do_unlink=True)
+        
+        # Clean up orphaned data blocks
+        for mesh in bpy.data.meshes:
+            if mesh.users == 0:
+                bpy.data.meshes.remove(mesh)
+        
+        for material in bpy.data.materials:
+            if material.users == 0:
+                bpy.data.materials.remove(material)
+        
+        for light_data in bpy.data.lights:
+            if light_data.users == 0:
+                bpy.data.lights.remove(light_data)
+
+
+class BlenderSegmaskRenderer:
+    def __init__(self): 
+        """
+        Initialize the Blender cuboid renderer.
+        
+        Args:
+            img_dim (int): Image dimensions (square)
+            render_engine (str): Blender render engine ('EEVEE' or 'CYCLES')
+            num_lights (int): Number of lights to add
+            max_tries (int): Maximum tries for placement
+        """
+        self.img_dim = 1024 
+        self.render_engine = "BLENDER_WORKBENCH"
+        self.blender_grid_dims = scales
+
+        self.radius = 6.0 
+        self.center = -6.0 
+        
+        # Scene references
+        self.context = None
+        self.scene = None
+        self.camera = None
+        self.render = None
+
+        # Setup the scene
+        self.setup_scene()
+
+        
+    def setup_scene(self):
+        """
+        Setup the basic Blender scene with camera, lighting, and render settings.
+        
+        Args:
+            camera_data (dict): Camera configuration containing elevation, lens, global_scale, etc.
+        """
+        # Get all objects in the scene
+        objects_to_remove = []
+        
+        for obj in bpy.data.objects:
+            # Remove default cube, plane, camera, and lights
+            if obj.type in {'MESH', 'LIGHT', 'CAMERA'}:
+                objects_to_remove.append(obj)
+        
+        # Delete the objects
+        for obj in objects_to_remove:
+            bpy.data.objects.remove(obj, do_unlink=True)
+        
+        # Also clear orphaned data
+        for mesh in bpy.data.meshes:
+            if mesh.users == 0:
+                bpy.data.meshes.remove(mesh)
+        
+        for light in bpy.data.lights:
+            if light.users == 0:
+                bpy.data.lights.remove(light)
+        
+        for camera in bpy.data.cameras:
+            if camera.users == 0:
+                bpy.data.cameras.remove(camera)
+
+        bpy.context.scene.world = None 
+
+        # Initialize Blender scene
+        # bpy.ops.wm.read_factory_settings(use_empty=True)
+        self.context = bpy.context 
+        self.scene = self.context.scene 
+        self.render = self.scene.render 
+        
+        # Set render engine and resolution
+        self.render.engine = self.render_engine  
+        self.context.scene.render.resolution_x = self.img_dim  
+        self.context.scene.render.resolution_y = self.img_dim  
+        self.context.scene.render.resolution_percentage = 100  
+
+        # Setup compositing nodes
+        self._setup_compositing()
+        
+        
+    def _setup_compositing(self):
+        """Setup Blender compositing nodes for depth and RGB output."""
+        self.context.scene.use_nodes = True
+        tree = self.context.scene.node_tree
+        links = tree.links
+
+        self.context.scene.render.use_compositing = True 
+        self.context.view_layer.use_pass_z = True
+        
+        # clear default nodes
+        for n in tree.nodes:
+            tree.nodes.remove(n)
+        
+        # create input render layer node
+        rl = tree.nodes.new('CompositorNodeRLayers')      
+
+        map_node = tree.nodes.new(type="CompositorNodeMapValue")
+        map_node.size = [0.05]
+        map_node.use_min = True
+        map_node.min = [0]
+        map_node.use_max = True
+        map_node.max = [65336]
+        links.new(rl.outputs[2], map_node.inputs[0])
+
+        invert = tree.nodes.new(type="CompositorNodeInvert")
+        links.new(map_node.outputs[0], invert.inputs[1])
+        
+        # create output node
+        v = tree.nodes.new('CompositorNodeViewer')   
+        v.use_alpha = True 
+
+        # create a file output node and set the path
+        fileOutput = tree.nodes.new(type="CompositorNodeOutputFile")
+        fileOutput.base_path = "."
+        links.new(invert.outputs[0], fileOutput.inputs[0])
+
+        # Links
+        links.new(rl.outputs[0], v.inputs[0])  # link Image to Viewer Image RGB
+        links.new(rl.outputs['Depth'], v.inputs[1])  # link Render Z to Viewer Image Alpha
+
+        # Update scene to apply changes
+        self.context.view_layer.update() 
+
+
+    def _setup_camera_cv(self, camera_data):
+        """Setup camera position and orientation."""
+        reset_cameras(self.scene) 
+        self.camera = self.scene.objects["Camera"] 
+        
+        elevation = camera_data["camera_elevation"] 
+        tan_elevation = np.tan(elevation) 
+        cos_elevation = np.cos(elevation) 
+        sin_elevation = np.sin(elevation) 
+
+        radius = self.radius 
+        center = self.center 
+        
+        self.camera.location = mathutils.Vector((radius * cos_elevation + center, 0, radius * sin_elevation))  
+        direction = mathutils.Vector((-1, 0, -tan_elevation)) 
+        self.context.scene.camera = self.camera 
+        rot_quat = direction.to_track_quat("-Z", "Y") 
+        self.camera.rotation_euler = rot_quat.to_euler() 
+        self.camera.data.lens = camera_data["lens"]  
+        
+    def _create_cuboid_objects(self, subjects_data):
+        """Create primitive cuboid objects for all subjects."""
+        for subject_idx, subject_data in enumerate(subjects_data): 
+            x = subject_data["x"][0] 
+            y = subject_data["y"][0] 
+            z = subject_data["z"][0] 
+            rgb_color = map_point_to_rgb(x, y, z) 
+            _, prim_obj = get_primitive_object(rgb_color) 
+            prim_obj.location = np.array([100, 0, 0]) 
+            subject_data["prim_obj"] = prim_obj
+
+    def _place_objects(self, subjects_data, camera_data):
+        """Place objects in the scene according to their data."""
+        global_scale = camera_data["global_scale"] 
+        
+        for subject_data in subjects_data: 
+            x = subject_data["x"][0]  
+            y = subject_data["y"][0]  
+            z = global_scale * subject_data["dims"][2] / 2.0 + subject_data["z"][0]  
+            subject_data["prim_obj"].location = np.array([x, y, z])  
+            subject_data["prim_obj"].scale = global_scale * np.array(subject_data["dims"]) * 2.0 
+            subject_data["prim_obj"].rotation_euler[2] = subject_data["azimuth"][0]  
+
+    def render_cv(self, subjects_data, camera_data, num_samples=1):
+        """
+        Main render method that takes subjects data and renders the scene.
+        
+        Args:
+            subjects_data (list): List of subject dictionaries containing position, dims, etc.
+            camera_data (dict): Camera configuration
+            num_samples (int): Number of samples to render (currently only supports 1)
+            output_path (str): Path to save the rendered image
+            
+        Returns:
+            None
+        """
+        # Setup camera
+        center = (-6.0, 0.0, 0.0) 
+        radius = 6.0 
+
+        for subject_data in subjects_data: 
+            subject_data["azimuth"][0] = np.deg2rad(subject_data["azimuth"][0]) 
+            subject_data["x"][0] = subject_data["x"][0] + center[0] 
+            subject_data["y"][0] = subject_data["y"][0] + center[1] 
+            subject_data["z"][0] = subject_data["z"][0] + center[2] 
+
+        print(f"in segmask render, {subjects_data = }") 
+
+        self._setup_camera_cv(camera_data)
+        
+        assert num_samples == 1, "for now, only implemented for a single sample"  
+        assert "global_scale" in camera_data.keys(), "global_scale must be set" 
+        
+        # Create primitive objects for subjects
+        self._create_cuboid_objects(subjects_data)
+
+        def make_segmask(image): 
+            alpha = image[:, :, 3] 
+            _, mask = cv2.threshold(alpha, 0, 255, cv2.THRESH_BINARY) 
+            return mask  
+
+
+        for subject_idx, subject_data in enumerate(subjects_data): 
+            # Place objects in scene
+            self._place_objects([subject_data], camera_data) 
+        
+            # Perform rendering
+            print(f"SUCCESS, rendering...")
+            self.context.scene.render.filepath = "tmp.png"
+            self.context.scene.render.image_settings.file_format = "PNG" 
+            bpy.ops.render.render(write_still=True) 
+            img = cv2.imread("tmp.png", cv2.IMREAD_UNCHANGED) 
+            segmask = make_segmask(img)  
+            print(f"{segmask.shape = }") 
+            cv2.imwrite(f"{str(subject_idx).zfill(3)}_segmask_cv.png", segmask) 
+            print(f"saved {str(subject_idx).zfill(3)}_segmask_cv.png") 
+
+            subject_data["prim_obj"].location = np.array([100, 0, 0]) # move out of view 
+        
+        self.cleanup() 
+
+
+    def cleanup(self):
+        """Clean up the scene for next render."""
+        # Remove all lights
+        remove_all_lights()
+        
+        # Remove all other objects (meshes, empties, etc.)
+        objects_to_remove = [obj for obj in bpy.data.objects]
+        
+        for obj in objects_to_remove:
+            bpy.data.objects.remove(obj, do_unlink=True)
+        
+        # Clean up orphaned data blocks
+        for mesh in bpy.data.meshes:
+            if mesh.users == 0:
+                bpy.data.meshes.remove(mesh)
+        
+        for material in bpy.data.materials:
+            if material.users == 0:
+                bpy.data.materials.remove(material)
+        
+        for light_data in bpy.data.lights:
+            if light_data.users == 0:
+                bpy.data.lights.remove(light_data)
+        
+
+        
+# Update the main execution
+if __name__ == '__main__':
+    subjects_data = [
+        {
+            "name": "sedan", 
+            "x": [-5.0],
+            "y": [0.0], 
+            "dims": [1.0, 2.0, 1.5],
+            "azimuth": [0.0]
+        }, 
+    ]
+    camera_data = {
+        "camera_elevation": np.arctan(0.45), 
+        "lens": 70,
+        "global_scale": 1.0
+    }
+    
+    # Create renderer instance
+    renderer = BlenderCuboidRenderer(
+        img_dim=1024,
+        render_engine='EEVEE',
+        num_lights=1,
+    )
+    
+    # Render the scene
+    renderer.render(
+        subjects_data=subjects_data, 
+        camera_data=camera_data, 
+        num_samples=1,
+        output_path="main.jpg"
+    )

gradio_app/blender_server.py

@@ -0,0 +1,149 @@
+import os
+import sys
+import tempfile
+import shutil
+import base64
+import io
+from PIL import Image
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Dict, Any
+import uvicorn
+import argparse
+
+# Import BlenderCuboidRenderer
+from blender_backend import BlenderCuboidRenderer
+
+class RenderRequest(BaseModel):
+    subjects_data: List[Dict[str, Any]]
+    camera_data: Dict[str, Any]
+    num_samples: int = 1
+
+class RenderResponse(BaseModel):
+    success: bool
+    image_base64: str = None
+    error_message: str = None
+
+class BlenderRenderServer:
+    def __init__(self, render_mode: str):
+        """
+        Initialize the Blender render server.
+        
+        Args:
+            render_mode (str): Either 'cv' for camera view or 'bev' for bird's eye view
+        """
+        self.render_mode = render_mode
+        if self.render_mode == "cv":
+            self.renderer = BlenderCuboidRenderer("BLENDER_EEVEE_NEXT")
+        elif self.render_mode == "final":
+            self.renderer = BlenderCuboidRenderer("CYCLES")
+        elif self.render_mode == "paper":
+            self.renderer = BlenderCuboidRenderer("CYCLES")
+        
+    def process_render_request(self, request: RenderRequest) -> RenderResponse:
+        """Process a render request and return the result."""
+        # Create temporary directory for this render
+        output_path = os.path.join(f"{self.render_mode}_render.jpg")
+        
+        # Convert subjects_data format if needed
+        converted_subjects_data = self._convert_subjects_data(request.subjects_data)
+        
+        # Add required camera_data fields
+        camera_data = request.camera_data.copy()
+        camera_data["global_scale"] = camera_data.get("global_scale", 1.0)
+        
+        # Perform the render based on mode
+        if self.render_mode == "cv":
+            self.renderer.render_cv(
+                subjects_data=converted_subjects_data,
+                camera_data=camera_data,
+                num_samples=request.num_samples,
+                output_path=output_path
+            )
+        elif self.render_mode == "final":
+            self.renderer.render_final_representation(
+                subjects_data=converted_subjects_data,
+                camera_data=camera_data,
+                num_samples=request.num_samples,
+                output_path=output_path
+            )
+        elif self.render_mode == "paper":
+            self.renderer.render_paper_figure(
+                subjects_data=converted_subjects_data,
+                camera_data=camera_data,
+                num_samples=request.num_samples,
+                output_path=output_path
+            )
+        else:
+            raise ValueError(f"Invalid render mode: {self.render_mode}")
+        
+        # Read and encode the rendered image
+        if os.path.exists(output_path):
+            with open(output_path, "rb") as img_file:
+                img_data = img_file.read()
+                img_base64 = base64.b64encode(img_data).decode('utf-8')
+            
+            return RenderResponse(success=True, image_base64=img_base64)
+        else:
+            return RenderResponse(
+                success=False, 
+                error_message="Render output file not found"
+            )
+                
+    def _convert_subjects_data(self, subjects_data: List[Dict]) -> List[Dict]:
+        """Convert subjects data to the format expected by BlenderCuboidRenderer."""
+        converted = []
+        
+        for subject in subjects_data:
+            # Convert to the expected format with lists for x, y, azimuth
+            converted_subject = {
+                "name": subject.get("subject_name", "cuboid"),
+                "x": [subject["x"]],
+                "y": [subject["y"]],
+                "z": [subject["z"]],
+                "dims": [subject["width"], subject["depth"], subject["height"]],
+                "azimuth": [subject["azimuth"]]
+            }
+            converted.append(converted_subject)
+            
+        return converted
+
+# Create FastAPI app
+app = FastAPI(title="Blender Render Server")
+
+# Global server instance
+server = None
+
+@app.on_event("startup")
+def startup_event():
+    global server
+    render_mode = os.environ.get("RENDER_MODE") 
+    server = BlenderRenderServer(render_mode)
+    print(f"Blender Render Server started in {render_mode.upper()} mode")
+
+@app.post("/render", response_model=RenderResponse)
+def render_scene(request: RenderRequest):
+    """Render a scene and return the result."""
+    if server is None:
+        raise HTTPException(status_code=500, detail="Server not initialized")
+    
+    return server.process_render_request(request)
+
+@app.get("/health")
+def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy", "render_mode": server.render_mode if server else "unknown"}
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Blender Render Server")
+    parser.add_argument("--mode", choices=["cv", "final", "paper"], required=True, 
+                       help="Render mode: cv for camera view, bev for bird's eye view")
+    parser.add_argument("--port", type=int, default=5001, help="Port to run server on")
+    parser.add_argument("--host", default="127.0.0.1", help="Host to bind server to")
+    
+    args = parser.parse_args()
+    
+    # Set environment variable for the startup event
+    os.environ["RENDER_MODE"] = args.mode
+    
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

gradio_app/blender_server_segmasks.py

@@ -0,0 +1,133 @@
+import os
+import sys
+import tempfile
+import shutil
+import base64
+import io
+from PIL import Image
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Dict, Any
+import uvicorn
+import argparse
+
+# Import BlenderSegmaskRenderer
+from blender_backend import BlenderSegmaskRenderer
+
+class SegmaskRenderRequest(BaseModel):
+    subjects_data: List[Dict[str, Any]]
+    camera_data: Dict[str, Any]
+    num_samples: int = 1
+
+class SegmaskRenderResponse(BaseModel):
+    success: bool
+    segmasks_base64: List[str] = None
+    error_message: str = None
+
+class BlenderSegmaskRenderServer:
+    def __init__(self):
+        """Initialize the Blender segmentation mask render server."""
+        self.renderer = BlenderSegmaskRenderer()
+        
+    def process_render_request(self, request: SegmaskRenderRequest) -> SegmaskRenderResponse:
+        """Process a segmentation mask render request and return the result."""
+        try:
+            # Create temporary directory for this render
+            # Convert subjects_data format if needed
+            converted_subjects_data = self._convert_subjects_data(request.subjects_data)
+            
+            # Add required camera_data fields
+            camera_data = request.camera_data.copy()
+            camera_data["global_scale"] = camera_data.get("global_scale", 1.0)
+            
+            # Perform the render
+            self.renderer.render_cv(
+                subjects_data=converted_subjects_data,
+                camera_data=camera_data,
+                num_samples=request.num_samples
+            )
+            
+            # Read and encode all segmentation masks in order
+            segmasks_base64 = []
+            num_subjects = len(converted_subjects_data)
+            
+            for subject_idx in range(num_subjects):
+                segmask_path = os.path.join(f"{str(subject_idx).zfill(3)}_segmask_cv.png")
+                
+                if os.path.exists(segmask_path):
+                    with open(segmask_path, "rb") as img_file:
+                        img_data = img_file.read()
+                        img_base64 = base64.b64encode(img_data).decode('utf-8')
+                        segmasks_base64.append(img_base64)
+                else:
+                    # Return error if any segmask is missing
+                    return SegmaskRenderResponse(
+                        success=False,
+                        error_message=f"Segmentation mask for subject {subject_idx} not found"
+                    )
+            
+            
+            return SegmaskRenderResponse(
+                success=True,
+                segmasks_base64=segmasks_base64
+            )
+                
+        except Exception as e:
+            # Change back to original directory on error
+            
+            return SegmaskRenderResponse(
+                success=False,
+                error_message=f"Segmentation mask render failed: {str(e)}"
+            )
+    
+    def _convert_subjects_data(self, subjects_data: List[Dict]) -> List[Dict]:
+        """Convert subjects data to the format expected by BlenderSegmaskRenderer."""
+        converted = []
+        
+        for subject in subjects_data:
+            # Convert to the expected format with lists for x, y, azimuth
+            converted_subject = {
+                "name": subject.get("subject_name", "cuboid"),
+                "x": [subject["x"]],
+                "y": [subject["y"]],
+                "z": [subject["z"]],
+                "dims": [subject["width"], subject["depth"], subject["height"]],
+                "azimuth": [subject["azimuth"]]
+            }
+            converted.append(converted_subject)
+            
+        return converted
+
+# Create FastAPI app
+app = FastAPI(title="Blender Segmentation Mask Render Server")
+
+# Global server instance
+server = None
+
+@app.on_event("startup")
+def startup_event():
+    global server
+    server = BlenderSegmaskRenderServer()
+    print("Blender Segmentation Mask Render Server started")
+
+@app.post("/render_segmasks", response_model=SegmaskRenderResponse)
+def render_segmasks(request: SegmaskRenderRequest):
+    """Render segmentation masks and return the results."""
+    if server is None:
+        raise HTTPException(status_code=500, detail="Server not initialized")
+    
+    return server.process_render_request(request)
+
+@app.get("/health")
+def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy", "type": "segmentation_mask_renderer"}
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Blender Segmentation Mask Render Server")
+    parser.add_argument("--port", type=int, default=5003, help="Port to run server on")
+    parser.add_argument("--host", default="127.0.0.1", help="Host to bind server to")
+    
+    args = parser.parse_args()
+    
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

gradio_app/infer_backend.py

@@ -0,0 +1,342 @@
+import os
+import json
+import torch
+from PIL import Image
+import numpy as np
+from typing import Optional, List, Tuple
+from transformers import CLIPTokenizer, T5TokenizerFast
+
+import sys 
+sys.path.append("..") 
+from train.src.pipeline import FluxPipeline
+from train.src.transformer_flux import FluxTransformer2DModel
+from train.src.lora_helper import set_single_lora, set_multi_lora, unset_lora
+from train.src.jsonl_datasets import make_train_dataset, collate_fn
+
+
+class InferenceArgs:
+    """Arguments configuration for inference dataset loading"""
+    def __init__(self, jsonl_path: str, pretrained_model_name: str):
+        # Basic paths
+        self.current_train_data_dir = jsonl_path
+        self.inference_embeds_dir = "/archive/vaibhav.agrawal/a-bev-of-the-latents/inference_embeds_flux2"
+        self.pretrained_model_name_or_path = pretrained_model_name
+        
+        # Column configurations
+        self.subject_column = None  # Set to None since we're using spatial
+        self.spatial_column = "cv"
+        self.target_column = "target"
+        self.caption_column = "PLACEHOLDER_prompts"
+        
+        # Size configurations
+        self.cond_size = 512
+        self.noise_size = 512
+        
+        # Other required parameters
+        self.revision = None
+        self.variant = None
+        self.max_sequence_length = 512
+
+
+class InferenceEngine:
+    """
+    Handles model loading and inference for the Gradio interface.
+    Pre-loads the base model and dynamically loads LoRA weights based on checkpoint selection.
+    """
+    
+    def __init__(self, base_model_path: str = "black-forest-labs/FLUX.1-dev", device: str = "cuda"):
+        """
+        Initialize the inference engine with base model.
+        
+        Args:
+            base_model_path: Path to the base FLUX model
+            device: Device to run inference on (default: "cuda")
+        """
+        self.device = device
+        self.base_model_path = base_model_path
+        self.current_lora_path = None
+        
+        print(f"Loading base model from {base_model_path}...")
+        
+        # Load pipeline and transformer
+        self.pipe = FluxPipeline.from_pretrained(
+            base_model_path, 
+            torch_dtype=torch.bfloat16, 
+            device=device
+        )
+        
+        transformer = FluxTransformer2DModel.from_pretrained(
+            base_model_path, 
+            subfolder="transformer",
+            torch_dtype=torch.bfloat16, 
+            device=device
+        )
+        
+        self.pipe.transformer = transformer
+        self.pipe.to(device)
+        
+        # Load tokenizers (same as in train.py and infer.ipynb)
+        print("Loading tokenizers...")
+        self.tokenizer_one = CLIPTokenizer.from_pretrained(
+            base_model_path,
+            subfolder="tokenizer",
+            revision=None,
+        )
+        self.tokenizer_two = T5TokenizerFast.from_pretrained(
+            base_model_path,
+            subfolder="tokenizer_2", 
+            revision=None,
+        )
+        self.tokenizers = [self.tokenizer_one, self.tokenizer_two]
+        
+        print("Base model and tokenizers loaded successfully!")
+    
+    def load_lora(self, checkpoint_name: str, lora_weights: List[float] = [1.0]):
+        """
+        Load LoRA weights for a specific checkpoint.
+        
+        Args:
+            checkpoint_name: Name of the checkpoint (e.g., "checkpoint_1")
+            lora_weights: Weights for the LoRA adaptation
+        """
+        # Construct LoRA path
+        lora_path = f"/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids/{checkpoint_name}/lora.safetensors"
+
+        print(f"\n\nGOT THE FOLLOWING LORA PATH: {lora_path}\n\n") 
+        
+        # Check if path exists
+        if not os.path.exists(lora_path):
+            raise FileNotFoundError(f"LoRA checkpoint not found at: {lora_path}")
+        
+        # Only reload if it's a different checkpoint
+        if self.current_lora_path != lora_path:
+            print(f"Loading LoRA weights from {lora_path}...")
+            set_single_lora(
+                self.pipe.transformer, 
+                lora_path, 
+                lora_weights=lora_weights,
+                cond_size=512
+            )
+            self.current_lora_path = lora_path
+            print(f"LoRA weights loaded successfully!")
+        else:
+            print(f"LoRA already loaded for {checkpoint_name}")
+    
+    def clear_cache(self):
+        """Clear attention processor cache"""
+        for name, attn_processor in self.pipe.transformer.attn_processors.items():
+            if hasattr(attn_processor, 'bank_kv'):
+                attn_processor.bank_kv.clear()
+    
+    def tensor_to_image_list(self, tensor):
+        """Convert normalized tensor to PIL Image list"""
+        if tensor is None:
+            return []
+        
+        images = []
+        for img_tensor in tensor:
+            # Denormalize from [-1, 1] to [0, 1]
+            img = (img_tensor.cpu().permute(1, 2, 0) * 0.5 + 0.5).clamp(0, 1).numpy()
+            # Convert to [0, 255] uint8
+            img = (img * 255.0).astype(np.uint8)
+            images.append(Image.fromarray(img))
+        
+        return images
+    
+    def run_inference(
+        self,
+        jsonl_path: str,
+        checkpoint_name: str,
+        height: int = 512,
+        width: int = 512,
+        seed: int = 42,
+        guidance_scale: float = 3.5,
+        num_inference_steps: int = 25,
+        max_sequence_length: int = 512
+    ) -> Tuple[bool, Optional[Image.Image], str]:
+        """
+        Run inference using data from JSONL file.
+        Uses the same data loading pipeline as training (make_train_dataset).
+        
+        Args:
+            jsonl_path: Path to the JSONL file containing inference data
+            checkpoint_name: Name of checkpoint to use
+            height: Output image height
+            width: Output image width
+            seed: Random seed for generation
+            guidance_scale: Guidance scale for diffusion
+            num_inference_steps: Number of denoising steps
+            max_sequence_length: Maximum sequence length for text encoding
+            
+        Returns:
+            Tuple of (success: bool, image: PIL.Image or None, message: str)
+        """
+        try:
+            # Load LoRA for selected checkpoint
+            self.load_lora(checkpoint_name)
+            
+            # Check if JSONL file exists
+            if not os.path.exists(jsonl_path):
+                return False, None, f"JSONL file not found at: {jsonl_path}"
+            
+            # Create inference arguments
+            inference_args = InferenceArgs(
+                jsonl_path=jsonl_path,
+                pretrained_model_name=self.base_model_path
+            )
+            
+            # Create dataset using the same pipeline as training
+            print("Creating inference dataset...")
+            inference_dataset = make_train_dataset(inference_args, self.tokenizers, accelerator=None, noise_size=512)
+            
+            # Create dataloader with batch_size=1
+            inference_dataloader = torch.utils.data.DataLoader(
+                inference_dataset,
+                batch_size=1,
+                shuffle=False,
+                collate_fn=collate_fn,
+                num_workers=0,
+            )
+            
+            # Get the first (and only) batch
+            batch = next(iter(inference_dataloader))
+            
+            # Extract data from batch
+            caption = batch["prompts"][0] if isinstance(batch["prompts"], list) else batch["prompts"]
+            call_ids = batch["call_ids"]
+            
+            print(f"\n{'='*60}")
+            print(f"Running inference with:")
+            print(f"  Checkpoint: {checkpoint_name}")
+            print(f"  Prompt: {caption}")
+            print(f"  Call IDs: {call_ids}")
+            print(f"  Height: {height}, Width: {width}")
+            print(f"  Seed: {seed}, Steps: {num_inference_steps}")
+            print(f"  Guidance Scale: {guidance_scale}")
+            print(f"{'='*60}\n")
+            
+            # Convert spatial condition tensors to PIL Images
+            spatial_imgs = self.tensor_to_image_list(batch["cond_pixel_values"])
+            
+            # Prepare cuboids segmentation masks
+            cuboids_segmasks = batch.get("cuboids_segmasks", None)
+            
+            # Prepare joint attention kwargs
+            joint_attention_kwargs = {
+                "call_ids": call_ids,
+                "cuboids_segmasks": cuboids_segmasks,
+            }
+            
+            print(f"Spatial images: {len(spatial_imgs)}")
+            print(f"{len(cuboids_segmasks) = }, {cuboids_segmasks[0].shape = }")
+            # print(f"Cuboids segmasks shape: {cuboids_segmasks.shape if cuboids_segmasks is not None else 'None'}")
+            cuboids_segmasks = torch.stack(cuboids_segmasks, dim=0) if cuboids_segmasks is not None else None 
+            
+            # Run inference
+            image = self.pipe(
+                prompt=caption,
+                height=int(height),
+                width=int(width),
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                max_sequence_length=max_sequence_length,
+                generator=torch.Generator("cpu").manual_seed(seed),
+                subject_images=[],  # No subject images for spatial conditioning
+                spatial_images=spatial_imgs,
+                cond_size=512,
+                **joint_attention_kwargs
+            ).images[0]
+            
+            # Clear cache
+            self.clear_cache()
+            torch.cuda.empty_cache()
+            
+            success_msg = f"✅ Successfully generated image using {checkpoint_name}"
+            print(f"\n{success_msg}\n")
+            
+            return True, image, success_msg
+            
+        except Exception as e:
+            error_msg = f"❌ Inference failed: {str(e)}"
+            print(f"\n{error_msg}\n")
+            import traceback
+            traceback.print_exc()
+            return False, None, error_msg
+
+
+# Global inference engine instance
+_inference_engine: Optional[InferenceEngine] = None
+
+
+def initialize_inference_engine(base_model_path: str = "black-forest-labs/FLUX.1-dev"):
+    """
+    Initialize the global inference engine.
+    Should be called once when the Gradio demo starts.
+    """
+    global _inference_engine
+    
+    if _inference_engine is None:
+        print("\n" + "="*60)
+        print("INITIALIZING INFERENCE ENGINE")
+        print("="*60 + "\n")
+        
+        _inference_engine = InferenceEngine(base_model_path=base_model_path)
+        
+        print("\n" + "="*60)
+        print("INFERENCE ENGINE READY")
+        print("="*60 + "\n")
+    
+    return _inference_engine
+
+
+def get_inference_engine() -> InferenceEngine:
+    """
+    Get the global inference engine instance.
+    Raises an error if not initialized.
+    """
+    global _inference_engine
+    
+    if _inference_engine is None:
+        raise RuntimeError(
+            "Inference engine not initialized. "
+            "Call initialize_inference_engine() first."
+        )
+    
+    return _inference_engine
+
+
+def run_inference_from_gradio(
+    checkpoint_name: str,
+    height: int = 512,
+    width: int = 512,
+    seed: int = 42,
+    guidance_scale: float = 3.5,
+    num_inference_steps: int = 25,
+    jsonl_path: str = "/archive/vaibhav.agrawal/a-bev-of-the-latents/gradio_files/cuboids.jsonl"
+) -> Tuple[bool, Optional[Image.Image], str]:
+    """
+    Wrapper function to run inference from Gradio interface.
+    
+    Args:
+        checkpoint_name: Name of checkpoint to use (from dropdown)
+        height: Output image height
+        width: Output image width
+        seed: Random seed
+        guidance_scale: Guidance scale
+        num_inference_steps: Number of denoising steps
+        jsonl_path: Path to JSONL file with inference data
+        
+    Returns:
+        Tuple of (success, generated_image, status_message)
+    """
+    engine = get_inference_engine()
+    
+    return engine.run_inference(
+        jsonl_path=jsonl_path,
+        checkpoint_name=checkpoint_name,
+        height=height,
+        width=width,
+        seed=seed,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps
+    )

gradio_app/launch_blender_backend.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+PORTS=(5001 5002 5003 5004)
+
+for port in "${PORTS[@]}"; do
+  PID=$(lsof -t -i ":$port")
+  if [ -n "$PID" ]; then
+    echo "Killing process $PID running on port $port..."
+    kill -9 "$PID"
+    echo "Process $PID killed."
+  else
+    echo "No process found running on port $port."
+  fi
+done
+
+# Start CV render server
+echo "Starting Camera View render server on port 5001..."
+python blender_server.py --mode cv --port 5001 &
+CV_PID=$!
+
+echo "Starting Camera View render server on port 5002..."
+python blender_server.py --mode final --port 5002 &
+FINAL_PID=$!
+
+
+# Start segmask render server
+echo "Starting Segmentation Mask render server on port 5003..."
+python3 blender_server_segmasks.py --port 5003 &
+SEGMASK_PID=$!
+
+echo "Starting Camera View render server on port 5004..."
+python blender_server.py --mode paper --port 5004 &
+PAPER_PID=$!
+
+echo "Render servers started!"
+echo "CV Server PID: $CV_PID (port 5001)"
+echo "Final (Cycles) Render Server PID: $FINAL_PID (port 5002)"
+echo "Segmentation Mask Server PID: $SEGMASK_PID (port 5003)"
+
+# Function to cleanup on exit
+cleanup() {
+    echo "Stopping render servers..."
+    kill $CV_PID $FINAL_PID $SEGMASK_PID 2>/dev/null
+    exit 0
+}
+
+# Set trap to cleanup on script exit
+trap cleanup SIGINT SIGTERM
+
+# Wait for both processes
+wait $CV_PID $FINAL_PID $SEGMASK_PID  

gradio_app/object_scales.py

@@ -0,0 +1,212 @@
+scales = {
+    "bear": 0.53,       # Unchanged
+    "bicycle": 0.4,      # Unchanged
+    "bugatti": 1.0,      # Unchanged
+    "bulldozer": 1.78,   # Unchanged
+    "bus": 2.67,         # Unchanged
+    "cat": 0.11,         # Unchanged
+    "chair": 0.18,       # Unchanged
+    "coupe": 1.0,        # Unchanged
+    "cow": 0.56,         # Unchanged
+    "crow": 0.09,        # CHANGED: Reduced from 0.11
+    "deer": 0.44,        # Unchanged
+    "dog": 0.22,         # Unchanged
+    "elephant": 1.22,    # Unchanged
+    "ferrari": 1.05,     # CHANGED: Increased from 1.0
+    "flamingo": 0.10,    # Unchanged
+    "fox": 0.22,         # Unchanged
+    "giraffe": 0.90,     # CHANGED: Reduced from 1.00
+    "goat": 0.33,        # Unchanged
+    "helicopter": 2.26,  # Unchanged
+    "hen": 0.09,         # Unchanged
+    "horse": 0.53,       # Unchanged
+    "jeep": 0.96,        # Unchanged
+    "kangaroo": 0.38,    # CHANGED: Increased from 0.33
+    "lamborghini": 1.0,  # Unchanged
+    "lion": 0.56,        # Unchanged
+    "mclaren": 1.0,      # Unchanged
+    "motorbike": 0.44,   # Unchanged
+    "office chair": 0.20,# Unchanged
+    "pickup truck": 1.22,# Unchanged
+    "pigeon": 0.067,     # Unchanged
+    "pig": 0.33,         # Unchanged
+    "rabbit": 0.11,      # Unchanged
+    "scooter": 0.4,      # Unchanged
+    "sedan": 1.0,        # Unchanged (Reference)
+    "sheep": 0.29,       # Unchanged
+    "shoe": 0.04,        # Unchanged
+    "sparrow": 0.033,    # Unchanged
+    "suv": 1.07,         # Unchanged
+    "table": 0.4,        # Unchanged
+    "teddy": 0.05,       # CHANGED: Reduced from 0.11
+    "tiger": 0.67,       # Unchanged
+    "tractor": 0.80,     # Unchanged
+    "van": 1.11,         # Unchanged
+    "vw beetle": 1.0,    # Unchanged
+    "wolf": 0.33,        # Unchanged
+    "man": 0.38,         # Unchanged
+    "zebra": 0.56        # Unchanged
+}
+
+tiny_assets = [
+    "sparrow",      # 0.033
+    "shoe",         # 0.04
+    "teddy",        # 0.05 (CHANGED)
+    "pigeon",       # 0.067
+    "hen",          # 0.09
+    "crow",         # 0.09 (CHANGED)
+    "flamingo",     # 0.10 (CHANGED - Moved from small)
+    "rabbit",       # 0.11
+    "cat",          # 0.11
+]
+
+
+small_assets = [
+    "chair",        # 0.18
+    "office chair", # 0.20
+    "dog",          # 0.22
+    "fox",          # 0.22
+    "sheep",        # 0.29
+    "goat",         # 0.33
+    "pig",          # 0.33
+    "wolf",         # 0.33
+    "man",          # 0.38 (CHANGED - Added to group)
+    "kangaroo",     # 0.38 (CHANGED)
+]
+
+
+medium_assets = [
+    "table",        # 0.4 (CHANGED - Moved from small)
+    "bicycle",      # 0.4
+    "scooter",      # 0.4
+    "deer",         # 0.44
+    "motorbike",    # 0.44
+    "bear",         # 0.53
+    "horse",        # 0.53
+    "cow",          # 0.56
+    "lion",         # 0.56
+    "zebra",        # 0.56
+    "tiger",        # 0.67
+    "tractor",      # 0.80
+    "giraffe",      # 0.90 (CHANGED)
+    "jeep",         # 0.96
+    "bugatti",      # 1.0
+    "coupe",        # 1.0
+    "lamborghini",  # 1.0
+    "mclaren",      # 1.0
+    "sedan",        # 1.0
+    "vw beetle",    # 1.0
+    "ferrari",      # 1.05 (CHANGED)
+    "suv",          # 1.07
+    "van",          # 1.11
+    "elephant",     # 1.22
+    "pickup truck", # 1.22
+    "bulldozer",    # 1.78
+    "helicopter",   # 2.26
+    "bus",          # 2.67
+]
+
+tiny_prompts = [
+    "a photo of PLACEHOLDER in a cozy birdhouse nestled in a green tree",
+    "a photo of PLACEHOLDER on a sandy beach near the water's edge with small shells",
+    "a photo of PLACEHOLDER amongst colorful wildflowers in a sunny meadow",
+    "a photo of PLACEHOLDER on a moss-covered log in a quiet forest",
+    "a photo of PLACEHOLDER near a small pond with lily pads floating",
+    "a photo of PLACEHOLDER on a window sill overlooking a rainy city street",
+    "a photo of PLACEHOLDER in a child's bedroom surrounded by other toys",
+    "a photo of PLACEHOLDER on a park bench with fallen leaves around",
+    "a photo of PLACEHOLDER by a small stream with smooth pebbles",
+    "a photo of PLACEHOLDER in a field of tall grass swaying gently",
+    "a photo of PLACEHOLDER on a wooden fence post in the countryside",
+    "a photo of PLACEHOLDER amongst blossoming spring flowers in a garden",
+    "a photo of PLACEHOLDER on a stack of old books in a library",
+    "a photo of PLACEHOLDER near a bird feeder in a winter garden",
+    "a photo of PLACEHOLDER on a picnic blanket in a sunny park",
+    "a photo of PLACEHOLDER on a kitchen counter near ripe fruit",
+    "a photo of PLACEHOLDER amongst autumn leaves on a forest floor",
+    "a photo of PLACEHOLDER on a rocky outcrop with a distant view",
+    "a photo of PLACEHOLDER near a puddle reflecting the sky",
+    "a photo of PLACEHOLDER in a patch of soft green moss",
+    "a photo of PLACEHOLDER on a weathered stone wall",
+    "a photo of PLACEHOLDER near a patch of blooming daisies",
+    "a photo of PLACEHOLDER on a sandy path through a garden",
+    "a photo of PLACEHOLDER near a watering can in a greenhouse",
+    "a photo of PLACEHOLDER amongst fallen pine needles in a forest",
+    "a photo of PLACEHOLDER on a small bridge over a gentle stream",
+    "a photo of PLACEHOLDER near a patch of colorful mushrooms"
+]
+
+small_prompts = [
+    "a photo of PLACEHOLDER in a sun-drenched greenhouse surrounded by various plants",
+    "a photo of PLACEHOLDER in a bustling city park with people walking by",
+    "a photo of PLACEHOLDER in a cozy library with tall bookshelves and soft lighting",
+    "a photo of PLACEHOLDER on a sandy dune near the ocean with gentle waves",
+    "a photo of PLACEHOLDER amongst tall reeds in a marshland area",
+    "a photo of PLACEHOLDER in a quiet forest clearing with sunlight filtering through trees",
+    "a photo of PLACEHOLDER on a grassy hill overlooking a small town",
+    "a photo of PLACEHOLDER near a flowing waterfall with mist in the air",
+    "a photo of PLACEHOLDER in a vibrant flower market with colorful blooms all around",
+    "a photo of PLACEHOLDER on a wooden dock extending into a still lake",
+    "a photo of PLACEHOLDER amongst rows of crops in a rural farmland",
+    "a photo of PLACEHOLDER in a historic town square with old buildings",
+    "a photo of PLACEHOLDER on a rocky beach with crashing waves in the distance",
+    "a photo of PLACEHOLDER amongst tall bamboo stalks in a serene grove",
+    "a photo of PLACEHOLDER in a snowy field with tracks visible in the snow",
+    "a photo of PLACEHOLDER on a paved walkway in a botanical garden",
+    "a photo of PLACEHOLDER near a campfire in a forest at night",
+    "a photo of PLACEHOLDER amongst colorful autumn foliage in a park",
+    "a photo of PLACEHOLDER on a stone path winding through a garden",
+    "a photo of PLACEHOLDER in a misty meadow with dew-covered grass",
+    "a photo of PLACEHOLDER on a wooden bridge crossing a small river",
+    "a photo of PLACEHOLDER amongst blooming lavender fields under a sunny sky",
+    "a photo of PLACEHOLDER in a quiet suburban backyard with green grass",
+    "a photo of PLACEHOLDER on a rocky hillside with sparse vegetation",
+    "a photo of PLACEHOLDER near a clear mountain stream with smooth stones",
+    "a photo of PLACEHOLDER amongst fallen leaves in a shaded woodland",
+    "a photo of PLACEHOLDER on a grassy bank beside a calm canal",
+    "a photo of PLACEHOLDER in a vineyard with rows of grapevines",
+    "a photo of PLACEHOLDER near a traditional wooden farmhouse"
+]
+
+medium_prompts = [
+    "a photo of PLACEHOLDER in a vast open plain with a dramatic sunset on the horizon",
+    "a photo of PLACEHOLDER on a winding mountain road with scenic views of valleys",
+    "a photo of PLACEHOLDER in a bustling harbor with various boats and ships",
+    "a photo of PLACEHOLDER in a dense pine forest with tall trees reaching the sky",
+    "a photo of PLACEHOLDER on a sandy beach with palm trees swaying in the breeze",
+    "a photo of PLACEHOLDER amongst rolling hills in a green countryside landscape",
+    "a photo of PLACEHOLDER in a vibrant city square with historic architecture",
+    "a photo of PLACEHOLDER in a train yard with multiple railway tracks", 
+    "a photo of PLACEHOLDER amongst tall redwood trees in an ancient forest",
+    "a photo of PLACEHOLDER in a sprawling parking lot outside a shopping mall", 
+    "a photo of PLACEHOLDER on a coastal highway with ocean views and cliffs", 
+    "a photo of PLACEHOLDER amongst golden wheat fields under a clear summer sky",
+    "a photo of PLACEHOLDER in a rocky canyon with sparse desert vegetation and blue sky above",
+    "a photo of PLACEHOLDER on a grassy plateau overlooking a vast landscape",
+    "a photo of PLACEHOLDER in a snowy mountain range with visible ski slopes",
+    "a photo of PLACEHOLDER on a paved highway stretching across an open landscape",
+    "a photo of PLACEHOLDER amongst lush vegetation in a tropical rainforest",
+    "a photo of PLACEHOLDER in a historic European city with ornate buildings",
+    "a photo of PLACEHOLDER in front of the Eiffel Tower at sunset",
+    "a photo of PLACEHOLDER amongst tall sunflowers in a field under a bright sun",
+    "a photo of PLACEHOLDER in a deep valley with steep forested sides",
+    "a photo of PLACEHOLDER on a rocky coastline with crashing waves and sea spray",
+    "a photo of PLACEHOLDER amongst vineyards on rolling hills under a sunny sky",
+    "a photo of PLACEHOLDER in a wide open desert with distant mesas and clear air",
+    "a photo of PLACEHOLDER amongst autumn-colored trees along a winding river",
+    "a photo of PLACEHOLDER in a bustling marketplace with various stalls and people",
+    "a photo of PLACEHOLDER on a racing circuit with banked turns and grandstands", 
+    "a photo of PLACEHOLDER amongst tall grasses in a savanna landscape", 
+]
+
+groups = {
+    "tiny": tiny_assets,
+    "small": small_assets,
+    "medium": medium_assets,
+}
+
+groups_prompts = {
+    "tiny": tiny_prompts,  
+    "small": small_prompts, 
+    "medium": medium_prompts, 
+}

gradio_app/visualize_server.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""
+Flask web server to visualize inference results for 2-subject cases.
+Port: 7023
+"""
+
+import os
+import json
+from pathlib import Path
+from flask import Flask, render_template, send_from_directory
+import base64
+
+app = Flask(__name__)
+
+# Paths
+DATASET_FILE = "/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/cuboids_segmentation.jsonl"
+DATASET_ROOT = "/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval"
+RESULTS_DIR = "/archive/vaibhav.agrawal/a-bev-of-the-latents/VAL/results/omini_seg_baseline_r2_epoch-0_checkpoint-20000"
+
+def load_2_subject_cases():
+    """Load all 2-subject cases from the dataset."""
+    cases = []
+    with open(DATASET_FILE, 'r') as f:
+        for idx, line in enumerate(f):
+            data = json.loads(line)
+            if len(data['subjects']) == 2:
+                cases.append({
+                    'dataset_index': idx,
+                    'subjects': data['subjects'],
+                    'prompt': data['prompt'],
+                    'target': data['target'],
+                    'cv': data['cv']
+                })
+    return cases
+
+# Load cases on startup
+TWO_SUBJECT_CASES = load_2_subject_cases()
+print(f"Loaded {len(TWO_SUBJECT_CASES)} 2-subject cases")
+
+def get_image_path(case, image_type):
+    """Get the path for a specific image type."""
+    if image_type == 'ground_truth':
+        return os.path.join(DATASET_ROOT, case['target'])
+    elif image_type == 'segmentation':
+        return os.path.join(DATASET_ROOT, case['cv'])
+    elif image_type == 'generated':
+        # Find the generated image in results
+        viz_dir = os.path.join(RESULTS_DIR, 'generated_images')
+        # Pattern: sample_{sample_idx:04d}_idx_{dataset_index}_seed_{seed}.jpg
+        # We need to find the file that matches the dataset_index
+        if os.path.exists(viz_dir):
+            for filename in os.listdir(viz_dir):
+                if f"_idx_{case['dataset_index']}_" in filename:
+                    return os.path.join(viz_dir, filename)
+    return None
+
+@app.route('/')
+def index():
+    """Main page showing the first 2-subject case."""
+    return show_case(0)
+
+@app.route('/case/<int:case_idx>')
+def show_case(case_idx):
+    """Display a specific case."""
+    if case_idx < 0 or case_idx >= len(TWO_SUBJECT_CASES):
+        return "Case not found", 404
+    
+    case = TWO_SUBJECT_CASES[case_idx]
+    
+    # Get image paths
+    gt_path = get_image_path(case, 'ground_truth')
+    seg_path = get_image_path(case, 'segmentation')
+    gen_path = get_image_path(case, 'generated')
+    
+    # Check if files exist
+    gt_exists = os.path.exists(gt_path) if gt_path else False
+    seg_exists = os.path.exists(seg_path) if seg_path else False
+    gen_exists = os.path.exists(gen_path) if gen_path else False
+    
+    return render_template('viewer.html',
+                         case_idx=case_idx,
+                         total_cases=len(TWO_SUBJECT_CASES),
+                         subjects=', '.join(case['subjects']),
+                         prompt=case['prompt'].replace('PLACEHOLDER', ', '.join(case['subjects'])),
+                         dataset_index=case['dataset_index'],
+                         gt_exists=gt_exists,
+                         seg_exists=seg_exists,
+                         gen_exists=gen_exists,
+                         prev_idx=case_idx - 1 if case_idx > 0 else None,
+                         next_idx=case_idx + 1 if case_idx < len(TWO_SUBJECT_CASES) - 1 else None)
+
+@app.route('/image/<int:case_idx>/<image_type>')
+def serve_image(case_idx, image_type):
+    """Serve the requested image."""
+    if case_idx < 0 or case_idx >= len(TWO_SUBJECT_CASES):
+        return "Case not found", 404
+    
+    case = TWO_SUBJECT_CASES[case_idx]
+    image_path = get_image_path(case, image_type)
+    
+    if image_path and os.path.exists(image_path):
+        directory = os.path.dirname(image_path)
+        filename = os.path.basename(image_path)
+        return send_from_directory(directory, filename)
+    else:
+        return "Image not found", 404
+
+if __name__ == '__main__':
+    # Create templates directory if it doesn't exist
+    os.makedirs('templates', exist_ok=True)
+    
+    # Run server on all interfaces (0.0.0.0) for remote access
+    print(f"Starting server on port 7023...")
+    print(f"Access at: http://<your-host-ip>:7023")
+    app.run(host='0.0.0.0', port=7023, debug=True)

train/src/jsonl_datasets.py

@@ -107,8 +107,7 @@ def make_train_dataset(args, tokenizer, accelerator, noise_size, only_realistic_
                 prompt_file_name = "space_prompt.pth" 
             else: 
                 prompt_file_name = "_".join(caption.split(" ")) + ".pth" 
-            if args.inference_embeds_dir is not None: 
-                assert osp.exists(osp.join(args.inference_embeds_dir, prompt_file_name)), f"Prompt embeddings for '{caption}' not found in {args.inference_embeds_dir}. Please precompute and save them." 
+            if args.inference_embeds_dir is not None and osp.exists(osp.join(args.inference_embeds_dir, prompt_file_name)):
                 prompt_embeds = torch.load(osp.join(args.inference_embeds_dir, prompt_file_name), map_location="cpu") 
                 pooled_prompt_embeds = prompt_embeds["pooled_prompt_embeds"] 
                 prompt_embeds = prompt_embeds["prompt_embeds"] 

train/src/pipeline.py

@@ -619,11 +619,7 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
             cond_size=512,
             call_ids=None, 
             cuboids_segmasks=None,
-            store_qk=None, 
-            store_qk_timesteps=None, 
     ):
-        assert not ((store_qk is None) ^ (store_qk_timesteps is None)), "Please provide both store_qk and store_qk_timesteps or neither of them." 
-
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor
         self.cond_size = cond_size
@@ -756,12 +752,6 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
                 if self.interrupt:
                     continue
 
-                store_qk_ = copy.deepcopy(store_qk) 
-                if (store_qk_ is not None) and (i not in store_qk_timesteps):
-                    store_qk_ = None 
-                elif store_qk_ is not None: 
-                    store_qk_ = osp.join(store_qk, f"step_{i}") 
-
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latents.shape[0]).to(latents.dtype)
                 noise_pred = self.transformer(
@@ -777,7 +767,6 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
                     return_dict=False,
                     call_ids=call_ids, 
                     cuboids_segmasks=cuboids_segmasks, 
-                    store_qk=store_qk_, 
                 )[0]
 
                 # compute the previous noisy sample x_t -> x_t-1