first commit

.gitattributes

@@ -0,0 +1,3 @@
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

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)

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"
+    )

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")

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")

failed_images.txt

@@ -0,0 +1,101 @@
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/testing/
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/bicycle__van__scooter__sedan/001/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/dog__fox/000/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/elephant__ferrari__bear__sedan/009/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/goat__man__dog/006/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/dog__man/008/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/hen__crow/001/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/bicycle__giraffe/006/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/table__lion/008/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/deer__tiger/009/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/pigeon__teddy__hen/008/segmap.png
+/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard_eval/semantic_segmaps_dilated_eroded/bulldozer__mclaren__tiger/005/segmap.png

infer_backend.py

@@ -0,0 +1,340 @@
+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
+
+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)
+            
+            # 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
+    )

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  

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, 
+}

requirements.txt

@@ -0,0 +1,20 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch
+torchvision
+torchaudio
+diffusers
+easydict==1.13
+einops==0.8.1
+peft==0.17.0
+pillow==11.0.0
+protobuf==5.29.3
+requests==2.32.3
+safetensors==0.5.2
+sentencepiece==0.2.0
+spaces==0.34.1
+transformers==4.49.0
+datasets
+wandb
+matplotlib 
+opencv-python 
+wandb 

train/default_config.yaml

@@ -0,0 +1,16 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+main_process_port: 14121
+downcast_bf16: 'no'
+gpu_ids: 1,
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 1               
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

train/group_subjects.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, 
+}

train/make_jsonl2_clip.py

@@ -0,0 +1,271 @@
+import os
+import os.path as osp
+import json
+from tqdm import tqdm 
+import torch
+import pickle
+from group_subjects_tabletop import groups, groups_prompts
+from train import * 
+
+PRETRAINED_MODEL_NAME_OR_PATH = "black-forest-labs/FLUX.1-dev" 
+
+def load_clip_evaluation_results(eval_dir, subjects_comb, img_idx, img_name):
+	"""
+	Load the CLIP similarity results for a specific image.
+	
+	Args:
+		eval_dir: Base directory containing CLIP evaluation results
+		subjects_comb: Subject combination directory name
+		img_idx: Image index directory name
+		img_name: Name of the image file
+	
+	Returns:
+		Minimum CLIP similarity score, or None if file not found
+	"""
+	# Construct the path to the pkl file
+	pkl_filename = img_name.replace(".jpg", ".pkl")
+	pkl_path = osp.join(eval_dir, subjects_comb, img_idx, pkl_filename)
+	
+	if not osp.exists(pkl_path):
+		return None
+	
+	try:
+		with open(pkl_path, 'rb') as f:
+			data = pickle.load(f)
+		
+		similarities = data.get('similarities', [])
+		
+		# Filter out any potential zero scores for non-existent subjects
+		valid_similarities = [s for s in similarities if s > 0.0]
+		
+		if valid_similarities:
+			return min(valid_similarities)
+		else:
+			return 0.0
+			
+	except Exception as e:
+		print(f"Warning: Could not process file {pkl_path}. Error: {e}")
+		return None
+
+
+def get_call_ids_from_placeholder_prompt_flux(prompt: str, tokenizer_three, 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_three.encode("<placeholder>", return_tensors="pt")[0][:-1].item()  
+	prompt_tokens_three = tokenizer_three.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_three.batch_decode(tokenizer_three.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 generate_cuboids_jsonl(data_dir, output_path, subject_names_embeds_flux, tokenizer_one, tokenizer_two, 
+						   clip_eval_dir=None, min_clip_similarity=0.26):
+	"""
+	Generate a JSONL file for cuboids dataset similar to pose.jsonl format.
+	
+	Args:
+		data_dir: Path to the images directory (same as BlenderFLUXSyntheticDataset data_dir)
+		output_path: Path where the cuboids.jsonl file should be saved
+		clip_eval_dir: Directory containing CLIP evaluation results (optional)
+		min_clip_similarity: Minimum CLIP similarity threshold for depth_flux images (default: 0.26)
+	"""
+	
+	# Create inverse groups mapping
+	inverse_groups = {}
+	for category in groups:
+		for subject in groups[category]:
+			assert subject not in inverse_groups
+			inverse_groups[subject] = category
+	
+	jsonl_entries = []
+	filtered_count = 0
+	total_depth_flux = 0
+
+	imgs_dir = osp.join(data_dir, "main_imgs") 
+	cuboids_dir = osp.join(data_dir, "cuboids_monochrome") 
+	
+	# Iterate over the dataset structure (same as BlenderFLUXSyntheticDataset)
+	subjects_combs = os.listdir(imgs_dir) 
+	import random 
+	random.shuffle(subjects_combs) 
+	for subjects_comb in tqdm(subjects_combs):
+		if len(subjects_comb.split("__")) > 4: 
+			continue 
+		if subjects_comb.startswith("_"): 
+			continue 
+		subjects_ = subjects_comb.split("__")
+		subjects = [" ".join(subject_.split("_")) for subject_ in subjects_]
+		if "bed" in subjects: 
+			continue 
+		
+		subjects_groups = [inverse_groups[subject] for subject in subjects]
+		PROMPTS = groups_prompts[subjects_groups[-2]]
+		
+		subjects_comb_dir = osp.join(imgs_dir, subjects_comb)
+
+		assert clip_eval_dir is not None 
+
+		for img_idx in os.listdir(subjects_comb_dir):
+			if not osp.isdir(osp.join(subjects_comb_dir, img_idx)): 
+				continue 
+
+			img_idx_dir = osp.join(subjects_comb_dir, img_idx)
+			
+			# Check if required files exist
+			main_img_path = osp.join(img_idx_dir, "main.jpg")
+			cuboids_path = osp.join(cuboids_dir, subjects_comb, img_idx, "cuboids.jpg")
+			pkl_path = osp.join(img_idx_dir, "main.pkl")
+			
+			assert osp.exists(main_img_path), f"Main image path {main_img_path} does not exist"  
+			assert osp.exists(cuboids_path), f"Cuboids path {cuboids_path} does not exist"  
+			assert osp.exists(pkl_path), f"PKL path {pkl_path} does not exist" 
+
+			# Get all image types (depth_flux and rendering)
+			img_names = os.listdir(img_idx_dir)
+			
+			# Process depth_flux images (prompt*.jpg)
+			depth_flux_imgs = [img_name for img_name in img_names 
+							 if img_name.endswith(".jpg") and img_name.find("prompt") != -1 and img_name.find("DEBUG") == -1]
+			
+			# Filter depth_flux images based on CLIP similarity if eval_dir is provided
+			if clip_eval_dir is not None:
+				filtered_depth_flux_imgs = []
+				for img_name in depth_flux_imgs:
+					total_depth_flux += 1
+					min_similarity = load_clip_evaluation_results(clip_eval_dir, subjects_comb, img_idx, img_name)
+					
+					if min_similarity is not None and min_similarity >= min_clip_similarity:
+						filtered_depth_flux_imgs.append(img_name)
+					else:
+						filtered_count += 1
+				
+				depth_flux_imgs = filtered_depth_flux_imgs
+			
+			all_imgs = depth_flux_imgs + ["main.jpg"] 
+			# all_imgs = ["main.jpg"] 
+			
+			for img_name in all_imgs:
+				img_path = osp.join(img_idx_dir, img_name)
+				
+				if img_name != "main.jpg":  
+					# Extract prompt index and get corresponding prompt
+					prompt_idx = int(img_name.replace("prompt", "").replace(".jpg", ""))
+					print(f"{prompt_idx = }, {subjects_groups[-1] = }, {subjects_comb = }")
+					prompt = PROMPTS[prompt_idx]
+				else: 
+					prompt = "a photo of PLACEHOLDER" 
+				
+				# Create placeholder text
+				placeholder_text = ""
+				for subject in subjects[:-1]:
+					placeholder_text = placeholder_text + f"<placeholder> {subject} and "
+				for subject in subjects[-1:]:
+					placeholder_text = placeholder_text + f"<placeholder> {subject}"
+				placeholder_text = placeholder_text.strip()
+
+				subjects_embeds = [] 
+				cuboids_segmasks_paths = [] 
+				segmasks_dir = osp.join(data_dir, "cuboids_segmasks_cv", subjects_comb, img_idx) 
+				assert osp.exists(segmasks_dir) 
+				segmask_names = sorted(os.listdir(segmasks_dir))  
+				for subject_idx, subject in enumerate(subjects):
+					subject_embed_path = osp.join(subject_names_embeds_flux, f"{subject.replace(' ', '_')}.pth")
+					assert osp.exists(subject_embed_path), f"Subject embed path {subject_embed_path} does not exist" 
+					subject_embed_obj = torch.load(subject_embed_path) 
+					subjects_embeds.append(subject_embed_obj) 
+					cuboid_segmask_path = osp.join(data_dir, "cuboids_segmasks_cv", subjects_comb, img_idx, segmask_names[subject_idx])
+					cuboid_segmask_path = osp.relpath(cuboid_segmask_path, osp.dirname(output_path)) 
+					# assert osp.exists(cuboid_segmask_path), f"Cuboid segmask path {cuboid_segmask_path} does not exist"
+					cuboids_segmasks_paths.append(cuboid_segmask_path) 
+				placeholder_prompt = prompt 
+				prompt = prompt.replace("PLACEHOLDER", placeholder_text) 
+				call_ids = get_call_ids_from_placeholder_prompt_flux(prompt, tokenizer_two, subjects, subjects_embeds, debug=True)
+				print(f"{call_ids = }")
+				
+				# Create relative paths from the output jsonl location
+				rel_cuboids_path = osp.relpath(cuboids_path, osp.dirname(output_path))
+				rel_img_path = osp.relpath(img_path, osp.dirname(output_path))
+
+				# Create JSONL entry
+				entry = {
+					"cv": rel_cuboids_path,
+					"PLACEHOLDER_prompts": placeholder_prompt,  
+					"target": rel_img_path,
+					"subjects": subjects, 
+					"cuboids_segmasks": cuboids_segmasks_paths, 
+					"call_ids": call_ids, 
+				}
+				jsonl_entries.append(entry)
+	
+	# Print filtering statistics
+	if clip_eval_dir is not None:
+		print(f"\n--- Filtering Statistics ---")
+		print(f"Total depth_flux images evaluated: {total_depth_flux}")
+		print(f"Images filtered out (min similarity < {min_clip_similarity}): {filtered_count}")
+		print(f"Images retained: {total_depth_flux - filtered_count}")
+		print(f"Retention rate: {((total_depth_flux - filtered_count) / total_depth_flux * 100):.2f}%")
+		print("---------------------------\n")
+	
+	# Write JSONL file
+	os.makedirs(osp.dirname(output_path), exist_ok=True)
+	with open(output_path, 'w') as f:
+		for entry in jsonl_entries:
+			f.write(json.dumps(entry) + '\n')
+	
+	print(f"Generated {len(jsonl_entries)} entries in {output_path}")
+
+if __name__ == "__main__":
+	# Configuration
+	data_dir = "/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv9"  # Replace with actual imgs_dir path
+	output_path = "/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv9/cuboids_monochrome.jsonl"
+	subjects_embeds_path = "/archive/vaibhav.agrawal/a-bev-of-the-latents/subject_names_embeds_flux"  # Path to subject embeddings JSON 
+	clip_eval_dir = "/archive/vaibhav.agrawal/a-bev-of-the-latents/clip_evaluation__datasetv9"  # CLIP evaluation results directory
+	min_clip_similarity = 0.26  # Minimum CLIP similarity threshold
+	rendered_imgs_prompt = "An image of PLACEHOLDER"  # Customize as needed
+	
+	# You can also accept command line arguments
+	import sys
+	if len(sys.argv) >= 2:
+		imgs_dir = sys.argv[1]
+	if len(sys.argv) >= 3:
+		output_path = sys.argv[2]
+	if len(sys.argv) >= 4:
+		rendered_imgs_prompt = sys.argv[3]
+
+	tokenizer_one = CLIPTokenizer.from_pretrained(
+		PRETRAINED_MODEL_NAME_OR_PATH, 
+		subfolder="tokenizer",
+		revision=None, 
+	)
+	tokenizer_two = T5TokenizerFast.from_pretrained(
+		PRETRAINED_MODEL_NAME_OR_PATH, 
+		subfolder="tokenizer_2",
+		revision=None, 
+	)
+
+	placeholder_token_str = ["<placeholder>"]
+	num_added_tokens = tokenizer_one.add_tokens(placeholder_token_str)   
+	assert num_added_tokens == 1 
+	num_added_tokens = tokenizer_two.add_tokens(placeholder_token_str)   
+	assert num_added_tokens == 1 
+	
+	generate_cuboids_jsonl(data_dir, output_path, subjects_embeds_path, tokenizer_two, tokenizer_two,
+						   clip_eval_dir=clip_eval_dir, min_clip_similarity=min_clip_similarity)

train/merge_jsonls.py

@@ -0,0 +1,181 @@
+import json
+import sys
+from pathlib import Path
+from typing import Set, Dict, Any
+
+
+def extract_subjects_comb(path: str) -> str:
+    """
+    Extract subjects_comb from a path.
+    Format: cuboids_monochrome/subjects_comb/img_idx/cuboids.jpg
+    """
+    path_parts = path.split('/')
+    
+    if 'cuboids_monochrome' in path_parts:
+        idx = path_parts.index('cuboids_monochrome')
+        if idx + 1 < len(path_parts):
+            return path_parts[idx + 1]
+    
+    return None
+
+
+def get_subjects_combs_from_jsonl(jsonl_path: str) -> Set[str]:
+    """Extract all unique subjects_comb from a JSONL file."""
+    subjects_combs = set()
+    
+    with open(jsonl_path, 'r') as f:
+        for line_num, line in enumerate(f, 1):
+            try:
+                entry = json.loads(line.strip())
+                cv_path = entry.get('cv', '')
+                subjects_comb = extract_subjects_comb(cv_path)
+                
+                if subjects_comb:
+                    subjects_combs.add(subjects_comb)
+                    
+            except json.JSONDecodeError as e:
+                print(f"Warning: Could not parse line {line_num} in {jsonl_path}: {e}")
+                continue
+    
+    return subjects_combs
+
+
+def get_common_keys(jsonl_path1: str, jsonl_path2: str) -> Set[str]:
+    """Get keys that are common across ALL entries in both JSONL files."""
+    keys1 = None
+    keys2 = None
+    
+    # Get keys from first file
+    with open(jsonl_path1, 'r') as f:
+        for line in f:
+            try:
+                entry = json.loads(line.strip())
+                current_keys = set(entry.keys())
+                if keys1 is None:
+                    keys1 = current_keys
+                else:
+                    keys1 = keys1.intersection(current_keys)
+            except json.JSONDecodeError:
+                continue
+    
+    # Get keys from second file
+    with open(jsonl_path2, 'r') as f:
+        for line in f:
+            try:
+                entry = json.loads(line.strip())
+                current_keys = set(entry.keys())
+                if keys2 is None:
+                    keys2 = current_keys
+                else:
+                    keys2 = keys2.intersection(current_keys)
+            except json.JSONDecodeError:
+                continue
+    
+    # Return intersection of keys from both files
+    if keys1 is None or keys2 is None:
+        return set()
+    
+    common_keys = keys1.intersection(keys2)
+    return common_keys
+
+
+def merge_jsonls(jsonl_path1: str, jsonl_path2: str, output_path: str):
+    """
+    Merge two JSONL files, ensuring:
+    1. No overlapping subjects_comb (assertion)
+    2. Only common keys are included in output
+    3. All entries are concatenated
+    """
+    print(f"Checking for overlapping subjects_comb...")
+    combs1 = get_subjects_combs_from_jsonl(jsonl_path1)
+    combs2 = get_subjects_combs_from_jsonl(jsonl_path2)
+    
+    overlap = combs1.intersection(combs2)
+    
+    # Assert no overlap
+    assert len(overlap) == 0, (
+        f"ERROR: Found {len(overlap)} overlapping subjects_comb between files!\n"
+        f"Overlapping subjects_comb: {sorted(overlap)}"
+    )
+    
+    print(f"✓ No overlapping subjects_comb found")
+    print(f"  File 1: {len(combs1)} unique subjects_comb")
+    print(f"  File 2: {len(combs2)} unique subjects_comb")
+    
+    # Get common keys
+    print(f"\nFinding common keys...")
+    common_keys = get_common_keys(jsonl_path1, jsonl_path2)
+    print(f"{common_keys = }")
+    
+    assert len(common_keys) > 0, "ERROR: No common keys found between the two JSONL files!"
+    
+    print(f"✓ Found {len(common_keys)} common keys: {sorted(common_keys)}")
+    
+    # Merge files
+    print(f"\nMerging files to {output_path}...")
+    total_entries = 0
+    
+    with open(output_path, 'w') as out_f:
+        # Write entries from first file
+        with open(jsonl_path1, 'r') as f1:
+            for line_num, line in enumerate(f1, 1):
+                try:
+                    entry = json.loads(line.strip())
+                    # Keep only common keys
+                    filtered_entry = {k: entry[k] for k in common_keys if k in entry}
+                    out_f.write(json.dumps(filtered_entry) + '\n')
+                    total_entries += 1
+                except json.JSONDecodeError as e:
+                    print(f"Warning: Could not parse line {line_num} in {jsonl_path1}: {e}")
+                    continue
+        
+        # Write entries from second file
+        with open(jsonl_path2, 'r') as f2:
+            for line_num, line in enumerate(f2, 1):
+                try:
+                    entry = json.loads(line.strip())
+                    # Keep only common keys
+                    filtered_entry = {k: entry[k] for k in common_keys if k in entry}
+                    out_f.write(json.dumps(filtered_entry) + '\n')
+                    total_entries += 1
+                except json.JSONDecodeError as e:
+                    print(f"Warning: Could not parse line {line_num} in {jsonl_path2}: {e}")
+                    continue
+    
+    print(f"✓ Merged {total_entries} entries to {output_path}")
+    print(f"\nMerge complete!")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 4:
+        print("Usage: python merge_two_jsonls.py <jsonl_file1> <jsonl_file2> <output_jsonl>")
+        print("\nExample:")
+        print("  python merge_two_jsonls.py dataset1/cuboids.jsonl dataset2/cuboids.jsonl merged_cuboids.jsonl")
+        sys.exit(1)
+    
+    jsonl_path1 = sys.argv[1]
+    jsonl_path2 = sys.argv[2]
+    output_path = sys.argv[3]
+    
+    # Validate input files
+    if not Path(jsonl_path1).exists():
+        print(f"Error: File not found: {jsonl_path1}")
+        sys.exit(1)
+    
+    if not Path(jsonl_path2).exists():
+        print(f"Error: File not found: {jsonl_path2}")
+        sys.exit(1)
+    
+    # Check if output file already exists
+    if Path(output_path).exists():
+        response = input(f"Warning: {output_path} already exists. Overwrite? (y/n): ")
+        if response.lower() != 'y':
+            print("Aborted.")
+            sys.exit(0)
+    
+    # Merge files
+    try:
+        merge_jsonls(jsonl_path1, jsonl_path2, output_path)
+    except AssertionError as e:
+        print(f"\n{e}")
+        sys.exit(1)

train/src/jsonl_datasets.py

@@ -0,0 +1,348 @@
+from PIL import Image
+from datasets import load_dataset
+from torchvision import transforms
+import random
+import torch.nn.functional as F 
+import torch
+import os 
+import os.path as osp 
+import cv2
+
+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
+
+Image.MAX_IMAGE_PIXELS = None
+
+def multiple_16(num: float):
+    return int(round(num / 16) * 16)
+
+def get_random_resolution(min_size=512, max_size=1280, multiple=16):
+    resolution = random.randint(min_size // multiple, max_size // multiple) * multiple
+    return resolution
+
+def load_image_safely(image_path, size):
+    try:
+        image = Image.open(image_path).convert("RGB")
+        return image
+    except Exception as e:
+        print("file error: "+image_path)
+        with open("failed_images.txt", "a") as f:
+            f.write(f"{image_path}\n")
+        return Image.new("RGB", (size, size), (255, 255, 255))
+    
+def make_train_dataset(args, tokenizer, accelerator):
+    if args.current_train_data_dir is not None:
+        print("load_data")
+        dataset = load_dataset('json', data_files=args.current_train_data_dir)
+
+    # Add index column to the dataset
+    dataset = dataset.map(lambda examples, indices: {**examples, 'index': indices}, with_indices=True, batched=True)
+
+    column_names = dataset["train"].column_names
+    
+    # 6. Get the column names for input/target.
+    target_column = args.target_column
+    if args.subject_column is not None:
+        subject_columns = args.subject_column.split(",")
+    if args.spatial_column is not None:
+        spatial_columns= args.spatial_column.split(",")
+    
+    size = args.cond_size
+    # by default the noise size would be randomly sampled from (512, 1024) 
+    # noise_size = get_random_resolution(max_size=args.noise_size) # maybe 768 or higher
+    noise_size = get_random_resolution(min_size=512, max_size=512) # maybe 768 or higher
+    # subject_cond_train_transforms = transforms.Compose(
+    #     [
+    #         transforms.Lambda(lambda img: img.resize((
+    #             multiple_16(size * img.size[0] / max(img.size)),
+    #             multiple_16(size * img.size[1] / max(img.size))
+    #         ), resample=Image.BILINEAR)),
+    #         transforms.RandomHorizontalFlip(p=0.7),
+    #         transforms.RandomRotation(degrees=20),
+    #         transforms.Lambda(lambda img: transforms.Pad(
+    #             padding=(
+    #                 int((size - img.size[0]) / 2),
+    #                 int((size - img.size[1]) / 2),
+    #                 int((size - img.size[0]) / 2),
+    #                 int((size - img.size[1]) / 2) 
+    #             ),
+    #             fill=0 
+    #         )(img)),
+    #         transforms.ToTensor(),
+    #         transforms.Normalize([0.5], [0.5]),
+    #     ]
+    # )
+    cond_train_transforms = transforms.Compose(
+        [
+            transforms.Resize((size, size), interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.CenterCrop((size, size)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ]
+    )
+    subject_cond_train_transforms = cond_train_transforms  
+
+    def train_transforms(image, noise_size):
+        train_transforms_ = transforms.Compose(
+            [
+                transforms.Lambda(lambda img: img.resize((
+                    multiple_16(noise_size * img.size[0] / max(img.size)),
+                    multiple_16(noise_size * img.size[1] / max(img.size))
+                ), resample=Image.BILINEAR)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+        transformed_image = train_transforms_(image)
+        return transformed_image
+    
+    def load_and_transform_cond_images(images):
+        transformed_images = [cond_train_transforms(image) for image in images]
+        concatenated_image = torch.cat(transformed_images, dim=1)
+        return concatenated_image
+    
+    def load_and_transform_subject_images(images):
+        transformed_images = [subject_cond_train_transforms(image) for image in images]
+        concatenated_image = torch.cat(transformed_images, dim=1)
+        return concatenated_image
+    
+    tokenizer_clip = tokenizer[0]
+    tokenizer_t5 = tokenizer[1]
+
+    def retrieve_prompt_embeds_from_disk(args, examples):
+        captions = []
+        for caption in examples["prompts"]:
+            if isinstance(caption, str):
+                if random.random() < 0.1:
+                    captions.append(" ")  # 将文本设为空
+                else:
+                    captions.append(caption)
+            elif isinstance(caption, list):
+                raise NotImplementedError("list of captions not supported yet")
+                # take a random caption if there are multiple
+                if random.random() < 0.1:
+                    captions.append(" ")
+                else:
+                    captions.append(random.choice(caption))
+            else:
+                raise ValueError(
+                    f"Caption column should contain either strings or lists of strings."
+                )
+
+        all_prompt_embeds = [] 
+        all_pooled_prompt_embeds = [] 
+        for caption in captions: 
+            if caption == " ": 
+                prompt_file_name = "space_prompt.pth" 
+            else: 
+                prompt_file_name = "_".join(caption.split(" ")) + ".pth" 
+            if 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"] 
+            else: 
+                # raise FileNotFoundError(f"Prompt embeddings for '{caption}' not found in {args.inference_embeds_dir}. Please precompute and save them.") 
+                prompt_embeds = torch.zeros((1, 77, 768))  # Placeholder tensor
+                pooled_prompt_embeds = torch.zeros((1, 768))  # Placeholder tensor
+            all_prompt_embeds.append(prompt_embeds.squeeze(0)) 
+            all_pooled_prompt_embeds.append(pooled_prompt_embeds.squeeze(0)) 
+        return all_prompt_embeds, all_pooled_prompt_embeds 
+
+
+    def tokenize_prompt_clip_t5(examples):
+        captions = []
+        for caption in examples["prompts"]:
+            if isinstance(caption, str):
+                if random.random() < 0.1:
+                    captions.append(" ")  # 将文本设为空
+                else:
+                    captions.append(caption)
+            elif isinstance(caption, list):
+                # take a random caption if there are multiple
+                if random.random() < 0.1:
+                    captions.append(" ")
+                else:
+                    captions.append(random.choice(caption))
+            else:
+                raise ValueError(
+                    f"Caption column should contain either strings or lists of strings."
+                )
+        text_inputs = tokenizer_clip(
+            captions,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids_1 = text_inputs.input_ids
+
+        text_inputs = tokenizer_t5(
+            captions,
+            padding="max_length",
+            max_length=512,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids_2 = text_inputs.input_ids
+        return text_input_ids_1, text_input_ids_2
+
+    def preprocess_train(examples):
+        _examples = {}
+        train_data_dir = osp.dirname(args.current_train_data_dir)
+        if args.subject_column is not None:
+            subject_images = [[load_image_safely(osp.join(train_data_dir, examples[column][i]), args.cond_size) for column in subject_columns] for i in range(len(examples[target_column]))]
+            _examples["subject_pixel_values"] = [load_and_transform_subject_images(subject) for subject in subject_images]
+        if args.spatial_column is not None:
+            # this now has two conditions 
+            spatial_images = [[load_image_safely(osp.join(train_data_dir, examples[column][i]), args.cond_size) for column in spatial_columns] for i in range(len(examples[target_column]))]
+            _examples["cond_pixel_values"] = [load_and_transform_cond_images(spatial) for spatial in spatial_images]
+        target_images = [load_image_safely(osp.join(train_data_dir, image_path), args.cond_size) for image_path in examples[target_column]]
+        _examples["pixel_values"] = [train_transforms(image, noise_size) for image in target_images]
+        _examples["PLACEHOLDER_prompts"] = examples["PLACEHOLDER_prompts"]
+        subjects = examples["subjects"] 
+        _examples["subjects"] = subjects 
+        subjects_ = ["_".join(subject) for subject in subjects] # get the subject names with "_" instead of space 
+        _examples["prompts"] = [] 
+        # getting the prompts by replacing the PLACEHOLDER in the prompt with the actual subject names  
+        for i in range(len(examples["subjects"])): 
+            # replace the subjects string in the PLACEHOLDER 
+            prompt = examples["PLACEHOLDER_prompts"][i] 
+            placeholder_string = " and ".join(subjects[i]) 
+            prompt = prompt.replace("PLACEHOLDER", placeholder_string) 
+            _examples["prompts"].append(prompt) 
+        _examples["prompt_embeds"], _examples["pooled_prompt_embeds"] = retrieve_prompt_embeds_from_disk(args, _examples) 
+        # gettin the z passed cuboids segmentation mask 
+        _examples["cuboids_segmasks"] = [] 
+
+        def generous_resize_batch(masks, new_h, new_w):
+            """
+            masks: torch.Tensor of shape (B, H, W), values in {0,1}
+            new_h, new_w: desired output size
+            """
+            B, H, W = masks.shape
+            masks = masks.unsqueeze(1).float()   # -> (B,1,H,W)
+
+            # Compute pooling kernel/stride
+            kh = H // new_h
+            kw = W // new_w
+            assert H % new_h == 0 and W % new_w == 0, \
+                "H and W must be divisible by new_h and new_w for exact block pooling"
+
+            out = F.max_pool2d(masks, kernel_size=(kh, kw), stride=(kh, kw))
+            return out.squeeze(1).byte()   # -> (B,new_h,new_w)
+
+        for i in range(len(_examples["subjects"])): 
+            segmasks_this_example = examples["cuboids_segmasks"][i] 
+            # the name of the segmask is of the format "segmask_00<subject_idx>__<depth_value>.png" 
+            depth_values_this_example = [osp.basename(segmasks_this_example[j]).split("__")[-1].split(".png")[0] for j in range(len(subjects[i]))] 
+            depth_values_this_example = torch.as_tensor([float(depth) for depth in depth_values_this_example]) 
+            assert len(segmasks_this_example) == len(subjects[i]), f"Number of segmentation masks {len(segmasks_this_example)} does not match number of subjects {len(subjects[i])} for example {i}" 
+            segmasks_this_example = [cv2.imread(osp.join(train_data_dir, segmasks_this_example[j]), cv2.IMREAD_UNCHANGED) for j in range(len(subjects[i]))] 
+            # segmasks_this_example = [cv2.resize(segmask, (32, 32), interpolation=cv2.INTER_NEAREST) for segmask in segmasks_this_example] 
+            segmasks_this_example = [torch.as_tensor(segmask, dtype=torch.uint8) for segmask in segmasks_this_example] 
+            segmasks_this_example = torch.stack(segmasks_this_example, dim=0) # (n_subjects, h, w) 
+            mask = segmasks_this_example > 128   
+            segmasks_this_example[mask] = 1 
+            segmasks_this_example[~mask] = 0 
+            segmasks_this_example = generous_resize_batch(segmasks_this_example, 32, 32) 
+            assert segmasks_this_example.shape == (len(subjects[i]), 32, 32), f"Segmentation masks shape {segmasks_this_example.shape} does not match expected shape {(len(subjects[i]), 32, 32)} for example {i}" 
+            # z_passed_segmask = do_z_pass(segmasks_this_example, depth_values_this_example) 
+            # print(f"{z_passed_segmask.shape = }, {segmasks_this_example.shape = }")
+            # _examples["cuboids_segmasks"].append(z_passed_segmask) 
+            _examples["cuboids_segmasks"].append(segmasks_this_example) 
+
+        _examples["token_ids_clip"], _examples["token_ids_t5"] = tokenize_prompt_clip_t5(_examples)
+        _examples["call_ids"] = examples["call_ids"] 
+        _examples["index"] = examples["index"] 
+
+        return _examples
+
+    if accelerator is not None:
+        with accelerator.main_process_first():
+            train_dataset = dataset["train"].with_transform(preprocess_train)
+    else:
+        train_dataset = dataset["train"].with_transform(preprocess_train)
+
+    return train_dataset
+
+
+def collate_fn(examples):
+    if examples[0].get("cond_pixel_values") is not None:
+        cond_pixel_values = torch.stack([example["cond_pixel_values"] for example in examples])
+        cond_pixel_values = cond_pixel_values.to(memory_format=torch.contiguous_format).float()
+    else:
+        cond_pixel_values = None
+    if examples[0].get("subject_pixel_values") is not None: 
+        subject_pixel_values = torch.stack([example["subject_pixel_values"] for example in examples])
+        subject_pixel_values = subject_pixel_values.to(memory_format=torch.contiguous_format).float()
+    else:
+        subject_pixel_values = None
+
+    target_pixel_values = torch.stack([example["pixel_values"] for example in examples])
+    target_pixel_values = target_pixel_values.to(memory_format=torch.contiguous_format).float()
+    token_ids_clip = torch.stack([torch.tensor(example["token_ids_clip"]) for example in examples])
+    token_ids_t5 = torch.stack([torch.tensor(example["token_ids_t5"]) for example in examples])
+    prompt_embeds = torch.stack([example["prompt_embeds"] for example in examples], dim=0)
+    pooled_prompt_embeds = torch.stack([example["pooled_prompt_embeds"] for example in examples], dim=0) 
+    prompts = [example["prompts"] for example in examples] 
+    call_ids = [example["call_ids"] for example in examples] 
+    cuboids_segmasks = [example["cuboids_segmasks"] for example in examples] if examples[0].get("cuboids_segmasks") is not None else None 
+    indices = [example["index"] for example in examples]  # Add this line
+
+    return {
+        "cond_pixel_values": cond_pixel_values,
+        "subject_pixel_values": subject_pixel_values,
+        "pixel_values": target_pixel_values,
+        "text_ids_1": token_ids_clip,
+        "text_ids_2": token_ids_t5,
+        "prompt_embeds": prompt_embeds,
+        "pooled_prompt_embeds": pooled_prompt_embeds,
+        "prompts": prompts, 
+        "call_ids": call_ids, 
+        "cuboids_segmasks": cuboids_segmasks, 
+        "index": indices, 
+    }

train/src/layers.py

@@ -0,0 +1,360 @@
+import inspect
+import math
+from typing import Callable, List, Optional, Tuple, Union
+from einops import rearrange
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch import Tensor
+from diffusers.models.attention_processor import Attention
+import os 
+import os.path as osp 
+import numpy as np 
+
+class LoRALinearLayer(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 4,
+        network_alpha: Optional[float] = None,
+        device: Optional[Union[torch.device, str]] = None,
+        dtype: Optional[torch.dtype] = None,
+        cond_width=512,
+        cond_height=512,
+        number=0,
+        n_loras=1
+    ):
+        super().__init__()
+        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
+        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
+        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
+        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
+        self.network_alpha = network_alpha
+        self.rank = rank
+        self.out_features = out_features
+        self.in_features = in_features
+
+        nn.init.normal_(self.down.weight, std=1 / rank)
+        nn.init.zeros_(self.up.weight)
+        
+        self.cond_height = cond_height
+        self.cond_width = cond_width
+        self.number = number
+        self.n_loras = n_loras
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_dtype = hidden_states.dtype
+        dtype = self.down.weight.dtype
+
+        #### img condition
+        batch_size = hidden_states.shape[0]
+        cond_size = self.cond_width // 8 * self.cond_height // 8 * 16 // 64
+        block_size =  hidden_states.shape[1] - cond_size * self.n_loras
+        shape = (batch_size, hidden_states.shape[1], 3072)
+        mask = torch.ones(shape, device=hidden_states.device, dtype=dtype) 
+        mask[:, :block_size+self.number*cond_size, :] = 0
+        mask[:, block_size+(self.number+1)*cond_size:, :] = 0
+        hidden_states = mask * hidden_states
+        ####
+        
+        down_hidden_states = self.down(hidden_states.to(dtype))
+        up_hidden_states = self.up(down_hidden_states)
+
+        if self.network_alpha is not None:
+            up_hidden_states *= self.network_alpha / self.rank
+
+        return up_hidden_states.to(orig_dtype)
+    
+
+class MultiSingleStreamBlockLoraProcessor(nn.Module):
+    def __init__(self, dim: int, ranks=[], lora_weights=[], network_alphas=[], device=None, dtype=None, cond_width=512, cond_height=512, n_loras=1):
+        super().__init__()
+        # Initialize a list to store the LoRA layers
+        self.n_loras = n_loras
+        self.cond_width = cond_width
+        self.cond_height = cond_height
+        
+        self.q_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.k_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.v_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.lora_weights = lora_weights
+        
+
+    def __call__(self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        use_cond = False,
+        call_ids = None, 
+        cuboids_segmasks: torch.Tensor = None, 
+        store_qk: Optional[str] = None, 
+    ) -> torch.FloatTensor:
+                
+        batch_size, seq_len, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        query = attn.to_q(hidden_states) 
+        key = attn.to_k(hidden_states) 
+        value = attn.to_v(hidden_states) 
+        
+        for i in range(self.n_loras):
+            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states)
+            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states)
+            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        cond_size = self.cond_width // 8 * self.cond_height // 8 * 16 // 64
+        block_size =  hidden_states.shape[1] - cond_size * self.n_loras
+        scaled_cond_size = cond_size
+        scaled_block_size = block_size
+        scaled_seq_len = query.shape[2]
+
+        num_cond_blocks = self.n_loras
+        mask = torch.ones((scaled_seq_len, scaled_seq_len), device=hidden_states.device)
+        # zero for all the 'allowed' connections 
+        mask[ :scaled_block_size, :] = 0  # First block_size row
+        for i in range(num_cond_blocks):
+            start = i * scaled_cond_size + scaled_block_size
+            end = (i + 1) * scaled_cond_size + scaled_block_size
+            mask[start:end, start:end] = 0  # Diagonal blocks
+
+        assert mask.shape[0] == scaled_block_size + num_cond_blocks*scaled_cond_size, f"{mask.shape = }, {scaled_block_size=}, {num_cond_blocks=}, {scaled_cond_size=}" 
+
+        if call_ids is not None: 
+            # repeat across batch size and heads 
+            mask = mask.unsqueeze(0).unsqueeze(0).repeat(len(call_ids), 1, 1, 1)  # (batch_size, num_heads, seq_len, seq_len) 
+            num_img_tokens = scaled_block_size - 512 
+            for batch_idx in range(len(call_ids)): 
+                call_ids_this_example = call_ids[batch_idx] 
+                for subject_idx, call_ids_this_subject in enumerate(call_ids_this_example): 
+                    # preparing the cuboid mask 
+                    cuboid_mask = cuboids_segmasks[batch_idx][subject_idx]  # (h, w) 
+                    # assert cuboid_mask.shape == (int(math.sqrt(num_img_tokens)), int(math.sqrt(num_img_tokens))), f"{cuboid_mask.shape=}, {num_img_tokens=}"  
+                    cuboid_mask = cuboid_mask.to(torch.bool)
+
+                    # assert scaled_block_size == scaled_cond_size + 512, f"{scaled_cond_size=}, {scaled_block_size=}"  
+                    for i in range(num_cond_blocks): 
+                        cuboid_mask = cuboids_segmasks[batch_idx][subject_idx]  # (h, w) 
+                        cuboid_mask = cuboid_mask.to(torch.bool) 
+                        # masking out the condition tokens -> text token attention map  
+                        mask_subset = mask[batch_idx, :, scaled_block_size + i*scaled_cond_size : scaled_block_size + (i+1)*scaled_cond_size, call_ids_this_subject] 
+                        # assert mask_subset.shape == (1, num_img_tokens, len(call_ids_this_subject)), f"{mask_subset.shape=}, {attn.heads=}, {num_img_tokens=}, {len(call_ids_this_subject)=}"   
+                        mask_subset[:, cuboid_mask.flatten()] = 0  # enable attention to cuboid regions 
+
+                        mask[batch_idx, :, scaled_block_size + i*scaled_cond_size : scaled_block_size + (i+1)*scaled_cond_size, call_ids_this_subject] = mask_subset 
+
+
+        mask = mask * -1e20
+        mask = mask.to(query.dtype)
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask)
+
+        if store_qk:  
+            attn_weights = query.detach().to(torch.float16) @ key.detach().to(torch.float16).transpose(-1, -2)  # (batch_size, num_heads, query_len, key_len)
+            attn_weights = attn_weights + mask 
+            attn_weights = torch.mean(torch.softmax(attn_weights, dim=-1), dim=1)  
+            attn_weights = attn_weights.cpu() 
+            os.makedirs(osp.dirname(store_qk), exist_ok=True) 
+            torch.save(attn_weights, store_qk + ".pth") 
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        cond_hidden_states = hidden_states[:, block_size:,:]
+        hidden_states = hidden_states[:, : block_size,:]
+
+        return hidden_states if not use_cond else (hidden_states, cond_hidden_states)
+
+
+class MultiDoubleStreamBlockLoraProcessor(nn.Module):
+    def __init__(self, dim: int, ranks=[], lora_weights=[], network_alphas=[], device=None, dtype=None, cond_width=512, cond_height=512, n_loras=1):
+        super().__init__()
+        
+        # Initialize a list to store the LoRA layers
+        self.n_loras = n_loras
+        self.cond_width = cond_width
+        self.cond_height = cond_height
+        self.q_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.k_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.v_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.proj_loras = nn.ModuleList([
+            LoRALinearLayer(dim, dim, ranks[i],network_alphas[i], device=device, dtype=dtype, cond_width=cond_width, cond_height=cond_height, number=i, n_loras=n_loras)
+            for i in range(n_loras)
+        ])
+        self.lora_weights = lora_weights
+
+
+    def __call__(self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        use_cond=False,
+        call_ids = None, 
+        cuboids_segmasks: torch.Tensor = None, 
+        store_qk: Optional[str] = None, 
+    ) -> torch.FloatTensor:
+        
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `context` projections.
+        inner_dim = 3072
+        head_dim = inner_dim // attn.heads
+        encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) 
+        encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+        encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+        encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+        encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+        encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+
+        if attn.norm_added_q is not None:
+            encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+        if attn.norm_added_k is not None:
+            encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+        
+        query = attn.to_q(hidden_states) 
+        key = attn.to_k(hidden_states) 
+        value = attn.to_v(hidden_states) 
+        for i in range(self.n_loras):
+            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states)
+            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states)
+            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+        
+        # attention
+        query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+        key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+        value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        cond_size = self.cond_width // 8 * self.cond_height // 8 * 16 // 64 
+        block_size =  hidden_states.shape[1] - cond_size * self.n_loras
+        scaled_cond_size = cond_size
+        scaled_seq_len = query.shape[2]
+        scaled_block_size = scaled_seq_len - cond_size * self.n_loras
+        
+        num_cond_blocks = self.n_loras
+        mask = torch.ones((scaled_seq_len, scaled_seq_len), device=hidden_states.device)
+        mask[ :scaled_block_size, :] = 0  # First block_size row
+        for i in range(num_cond_blocks):
+            start = i * scaled_cond_size + scaled_block_size
+            end = (i + 1) * scaled_cond_size + scaled_block_size
+            mask[start:end, start:end] = 0  # Diagonal blocks
+
+        assert mask.shape[0] == scaled_block_size + num_cond_blocks*scaled_cond_size, f"{mask.shape = }, {scaled_block_size=}, {num_cond_blocks=}, {scaled_cond_size=}" 
+
+        if call_ids is not None: 
+            # repeat across batch size and heads 
+            mask = mask.unsqueeze(0).unsqueeze(0).repeat(len(call_ids), 1, 1, 1)  # (batch_size, num_heads, seq_len, seq_len) 
+            num_img_tokens = scaled_block_size - 512 
+            for batch_idx in range(len(call_ids)): 
+                call_ids_this_example = call_ids[batch_idx] 
+                for subject_idx, call_ids_this_subject in enumerate(call_ids_this_example): 
+                    # preparing the cuboid mask 
+                    cuboid_mask = cuboids_segmasks[batch_idx][subject_idx]  # (h, w) 
+                    # assert cuboid_mask.shape == (int(math.sqrt(num_img_tokens)), int(math.sqrt(num_img_tokens))), f"{cuboid_mask.shape=}, {num_img_tokens=}, {scaled_block_size=}"  
+                    cuboid_mask = cuboid_mask.to(torch.bool)
+
+                    # assert scaled_block_size == scaled_cond_size + 512, f"{scaled_cond_size=}, {scaled_block_size=}"  
+                    for i in range(num_cond_blocks): 
+                        cuboid_mask = cuboids_segmasks[batch_idx][subject_idx]  # (h, w) 
+                        cuboid_mask = cuboid_mask.to(torch.bool) 
+                        # masking out the condition tokens -> text token attention map  
+                        mask_subset = mask[batch_idx, :, scaled_block_size + i*scaled_cond_size : scaled_block_size + (i+1)*scaled_cond_size, call_ids_this_subject] 
+                        # assert mask_subset.shape == (1, num_img_tokens, len(call_ids_this_subject)), f"{mask_subset.shape=}, {attn.heads=}, {num_img_tokens=}, {len(call_ids_this_subject)=}"   
+                        mask_subset[:, cuboid_mask.flatten()] = 0  # enable attention to cuboid regions 
+
+                        mask[batch_idx, :, scaled_block_size + i*scaled_cond_size : scaled_block_size + (i+1)*scaled_cond_size, call_ids_this_subject] = mask_subset 
+
+
+        mask = mask * -1e20
+        mask = mask.to(query.dtype)
+        
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask)
+
+        if store_qk:  
+            attn_weights = query.detach().to(torch.float16) @ key.detach().to(torch.float16).transpose(-1, -2)  # (batch_size, num_heads, query_len, key_len)
+            attn_weights = attn_weights + mask 
+            attn_weights = torch.mean(torch.softmax(attn_weights, dim=-1), dim=1)  
+            attn_weights = attn_weights.cpu() 
+            os.makedirs(osp.dirname(store_qk), exist_ok=True) 
+            torch.save(attn_weights, store_qk + ".pth") 
+
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+        
+        encoder_hidden_states, hidden_states = (
+            hidden_states[:, : encoder_hidden_states.shape[1]],
+            hidden_states[:, encoder_hidden_states.shape[1] :],
+        )
+
+        # Linear projection (with LoRA weight applied to each proj layer)
+        hidden_states = attn.to_out[0](hidden_states)
+        for i in range(self.n_loras):
+             hidden_states = hidden_states + self.lora_weights[i] * self.proj_loras[i](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+        
+        cond_hidden_states = hidden_states[:, block_size:,:]
+        hidden_states = hidden_states[:, :block_size,:]
+        
+        return (hidden_states, encoder_hidden_states, cond_hidden_states) if use_cond else (encoder_hidden_states, hidden_states)

train/src/lora_helper.py

@@ -0,0 +1,196 @@
+from diffusers.models.attention_processor import FluxAttnProcessor2_0
+from safetensors import safe_open
+import re
+import torch
+from .layers import MultiDoubleStreamBlockLoraProcessor, MultiSingleStreamBlockLoraProcessor
+
+device = "cuda"
+
+def load_safetensors(path):
+    tensors = {}
+    with safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            tensors[key] = f.get_tensor(key)
+    return tensors
+
+def get_lora_rank(checkpoint):
+    for k in checkpoint.keys():
+        if k.endswith(".down.weight"):
+            return checkpoint[k].shape[0]
+
+def load_checkpoint(local_path):
+    if local_path is not None:
+        if '.safetensors' in local_path:
+            print(f"Loading .safetensors checkpoint from {local_path}")
+            checkpoint = load_safetensors(local_path)
+        else:
+            print(f"Loading checkpoint from {local_path}")
+            checkpoint = torch.load(local_path, map_location='cpu')
+    return checkpoint
+
+def update_model_with_lora(checkpoint, lora_weights, transformer, cond_size):
+        number = len(lora_weights)
+        ranks = [get_lora_rank(checkpoint) for _ in range(number)]
+        lora_attn_procs = {}
+        double_blocks_idx = list(range(19))
+        single_blocks_idx = list(range(38))
+        for name, attn_processor in transformer.attn_processors.items():
+            match = re.search(r'\.(\d+)\.', name)
+            if match:
+                layer_index = int(match.group(1))
+            
+            if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
+                
+                lora_state_dicts = {}
+                for key, value in checkpoint.items():
+                    # Match based on the layer index in the key (assuming the key contains layer index)
+                    if re.search(r'\.(\d+)\.', key):
+                        checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+                        if checkpoint_layer_index == layer_index and key.startswith("transformer_blocks"):
+                            lora_state_dicts[key] = value
+                
+                lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
+                    dim=3072, ranks=ranks, network_alphas=ranks, lora_weights=lora_weights, device=device, dtype=torch.bfloat16, cond_width=cond_size, cond_height=cond_size, n_loras=number
+                )
+                
+                # Load the weights from the checkpoint dictionary into the corresponding layers
+                for n in range(number):
+                    lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].proj_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].proj_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].to(device)
+                
+            elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
+                
+                lora_state_dicts = {}
+                for key, value in checkpoint.items():
+                    # Match based on the layer index in the key (assuming the key contains layer index)
+                    if re.search(r'\.(\d+)\.', key):
+                        checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+                        if checkpoint_layer_index == layer_index and key.startswith("single_transformer_blocks"):
+                            lora_state_dicts[key] = value
+                
+                lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
+                    dim=3072, ranks=ranks, network_alphas=ranks, lora_weights=lora_weights, device=device, dtype=torch.bfloat16, cond_width=cond_size, cond_height=cond_size, n_loras=number
+                )
+                # Load the weights from the checkpoint dictionary into the corresponding layers
+                for n in range(number):
+                    lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
+                    lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
+                    lora_attn_procs[name].to(device)
+            else:
+                lora_attn_procs[name] = FluxAttnProcessor2_0()
+
+        transformer.set_attn_processor(lora_attn_procs)
+        
+
+def update_model_with_multi_lora(checkpoints, lora_weights, transformer, cond_size):
+        ck_number = len(checkpoints)
+        cond_lora_number = [len(ls) for ls in lora_weights]
+        cond_number = sum(cond_lora_number)
+        ranks = [get_lora_rank(checkpoint) for checkpoint in checkpoints]
+        multi_lora_weight = []
+        for ls in lora_weights:
+            for n in ls:
+                multi_lora_weight.append(n)
+        
+        lora_attn_procs = {}
+        double_blocks_idx = list(range(19))
+        single_blocks_idx = list(range(38))
+        for name, attn_processor in transformer.attn_processors.items():
+            match = re.search(r'\.(\d+)\.', name)
+            if match:
+                layer_index = int(match.group(1))
+            
+            if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
+                lora_state_dicts = [{} for _ in range(ck_number)]
+                for idx, checkpoint in enumerate(checkpoints):
+                    for key, value in checkpoint.items():
+                        # Match based on the layer index in the key (assuming the key contains layer index)
+                        if re.search(r'\.(\d+)\.', key):
+                            checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+                            if checkpoint_layer_index == layer_index and key.startswith("transformer_blocks"):
+                                lora_state_dicts[idx][key] = value
+                
+                lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
+                    dim=3072, ranks=ranks, network_alphas=ranks, lora_weights=multi_lora_weight, device=device, dtype=torch.bfloat16, cond_width=cond_size, cond_height=cond_size, n_loras=cond_number
+                )
+                
+                # Load the weights from the checkpoint dictionary into the corresponding layers
+                num = 0
+                for idx in range(ck_number):
+                    for n in range(cond_lora_number[idx]):
+                        lora_attn_procs[name].q_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.q_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].q_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.q_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].k_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.k_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].k_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.k_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].v_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.v_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].v_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.v_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].proj_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.proj_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].proj_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.proj_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].to(device)
+                        num += 1
+                
+            elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
+                
+                lora_state_dicts = [{} for _ in range(ck_number)]
+                for idx, checkpoint in enumerate(checkpoints):
+                    for key, value in checkpoint.items():
+                        # Match based on the layer index in the key (assuming the key contains layer index)
+                        if re.search(r'\.(\d+)\.', key):
+                            checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+                            if checkpoint_layer_index == layer_index and key.startswith("single_transformer_blocks"):
+                                lora_state_dicts[idx][key] = value
+                
+                lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
+                    dim=3072, ranks=ranks, network_alphas=ranks, lora_weights=multi_lora_weight, device=device, dtype=torch.bfloat16, cond_width=cond_size, cond_height=cond_size, n_loras=cond_number
+                )
+                # Load the weights from the checkpoint dictionary into the corresponding layers
+                num = 0
+                for idx in range(ck_number):
+                    for n in range(cond_lora_number[idx]):
+                        lora_attn_procs[name].q_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.q_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].q_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.q_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].k_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.k_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].k_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.k_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].v_loras[num].down.weight.data = lora_state_dicts[idx].get(f'{name}.v_loras.{n}.down.weight', None)
+                        lora_attn_procs[name].v_loras[num].up.weight.data = lora_state_dicts[idx].get(f'{name}.v_loras.{n}.up.weight', None)
+                        lora_attn_procs[name].to(device)
+                        num += 1
+
+            else:
+                lora_attn_procs[name] = FluxAttnProcessor2_0()
+
+        transformer.set_attn_processor(lora_attn_procs)
+
+
+def set_single_lora(transformer, local_path, lora_weights=[], cond_size=512):
+    checkpoint = load_checkpoint(local_path)
+    update_model_with_lora(checkpoint, lora_weights, transformer, cond_size)
+   
+def set_multi_lora(transformer, local_paths, lora_weights=[[]], cond_size=512):
+    checkpoints = [load_checkpoint(local_path) for local_path in local_paths]
+    update_model_with_multi_lora(checkpoints, lora_weights, transformer, cond_size)
+
+def unset_lora(transformer):
+    lora_attn_procs = {}
+    for name, attn_processor in transformer.attn_processors.items():
+        lora_attn_procs[name] = FluxAttnProcessor2_0()
+    transformer.set_attn_processor(lora_attn_procs)
+
+
+'''
+unset_lora(pipe.transformer)
+lora_path = "./lora.safetensors"
+lora_weights = [1, 1]
+set_lora(pipe.transformer, local_path=lora_path, lora_weights=lora_weights, cond_size=512)
+'''

train/src/pipeline.py

@@ -0,0 +1,824 @@
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+import copy 
+import os 
+import os.path as osp 
+
+from diffusers.image_processor import (VaeImageProcessor)
+from diffusers.loaders import FluxLoraLoaderMixin, FromSingleFileMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from torchvision.transforms.functional import pad
+from .transformer_flux import FluxTransformer2DModel
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+def calculate_shift(
+        image_seq_len,
+        base_seq_len: int = 256,
+        max_seq_len: int = 4096,
+        base_shift: float = 0.5,
+        max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+def prepare_latent_image_ids_2(height, width, device, dtype):
+    latent_image_ids = torch.zeros(height//2, width//2, 3, device=device, dtype=dtype)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height//2, device=device)[:, None]  # y坐标
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width//2, device=device)[None, :]   # x坐标
+    return latent_image_ids
+
+def prepare_latent_subject_ids(height, width, device, dtype):
+    latent_image_ids = torch.zeros(height // 2, width // 2, 3, device=device, dtype=dtype)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2, device=device)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2, device=device)[None, :]
+    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+    latent_image_ids = latent_image_ids.reshape(
+        latent_image_id_height * latent_image_id_width, latent_image_id_channels
+    )
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+def resize_position_encoding(batch_size, original_height, original_width, target_height, target_width, device, dtype):
+    latent_image_ids = prepare_latent_image_ids_2(original_height, original_width, device, dtype)
+    scale_h = original_height / target_height
+    scale_w = original_width / target_width
+    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+    latent_image_ids = latent_image_ids.reshape(
+        latent_image_id_height * latent_image_id_width, latent_image_id_channels
+    )
+    #spatial进行PE插值
+    latent_image_ids_resized = torch.zeros(target_height//2, target_width//2, 3, device=device, dtype=dtype)
+    for i in range(target_height//2):
+        for j in range(target_width//2):
+            latent_image_ids_resized[i, j, 1] = i*scale_h
+            latent_image_ids_resized[i, j, 2] = j*scale_w
+    cond_latent_image_id_height, cond_latent_image_id_width, cond_latent_image_id_channels = latent_image_ids_resized.shape
+    cond_latent_image_ids = latent_image_ids_resized.reshape(
+            cond_latent_image_id_height * cond_latent_image_id_width, cond_latent_image_id_channels
+        )
+    # latent_image_ids_ = torch.concat([latent_image_ids, cond_latent_image_ids], dim=0)
+    return latent_image_ids, cond_latent_image_ids #, latent_image_ids_
+    
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+        encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+        scheduler,
+        num_inference_steps: Optional[int] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        timesteps: Optional[List[int]] = None,
+        sigmas: Optional[List[float]] = None,
+        **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
+    r"""
+    The Flux pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+            self,
+            scheduler: FlowMatchEulerDiscreteScheduler,
+            vae: AutoencoderKL,
+            text_encoder: CLIPTextModel,
+            tokenizer: CLIPTokenizer,
+            text_encoder_2: T5EncoderModel,
+            tokenizer_2: T5TokenizerFast,
+            transformer: FluxTransformer2DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels)) if hasattr(self, "vae") and self.vae is not None else 16
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 64
+
+    def _get_t5_prompt_embeds(
+            self,
+            prompt: Union[str, List[str]] = None,
+            num_images_per_prompt: int = 1,
+            max_sequence_length: int = 512,
+            device: Optional[torch.device] = None,
+            dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1: -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+            self,
+            prompt: Union[str, List[str]],
+            num_images_per_prompt: int = 1,
+            device: Optional[torch.device] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1: -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    def encode_prompt(
+            self,
+            prompt: Union[str, List[str]],
+            prompt_2: Union[str, List[str]],
+            device: Optional[torch.device] = None,
+            num_images_per_prompt: int = 1,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            max_sequence_length: int = 512,
+            lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+        return prompt_embeds, pooled_prompt_embeds, text_ids
+
+    # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_inpaint.StableDiffusion3InpaintPipeline._encode_vae_image
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i: i + 1]), generator=generator[i])
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+        image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        return image_latents
+
+    def check_inputs(
+            self,
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=None,
+            pooled_prompt_embeds=None,
+            callback_on_step_end_tensor_inputs=None,
+            max_sequence_length=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+                k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    @staticmethod
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        height = height // vae_scale_factor
+        width = width // vae_scale_factor
+
+        latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height * 2, width * 2)
+
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def prepare_latents(
+            self,
+            batch_size,
+            num_channels_latents,
+            height,
+            width,
+            dtype,
+            device,
+            generator,
+            subject_image,
+            condition_image,
+            latents=None,
+            cond_number=1,
+            sub_number=1
+    ):
+        height_cond = 2 * (self.cond_size // self.vae_scale_factor)
+        width_cond = 2 * (self.cond_size // self.vae_scale_factor)
+        height = 2 * (int(height) // self.vae_scale_factor)  
+        width = 2 * (int(width) // self.vae_scale_factor)
+
+        shape = (batch_size, num_channels_latents, height, width)  # 1 16 106 80
+        noise_latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)  
+        noise_latents = self._pack_latents(noise_latents, batch_size, num_channels_latents, height, width)
+        noise_latent_image_ids, cond_latent_image_ids = resize_position_encoding(
+                batch_size,
+                height,
+                width,
+                height_cond,
+                width_cond,
+                device,
+                dtype,
+            )
+        
+        latents_to_concat = []  # 不包含 latents
+        latents_ids_to_concat = [noise_latent_image_ids]
+        
+        # subject
+        if subject_image is not None:
+            shape_subject = (batch_size, num_channels_latents, height_cond*sub_number, width_cond)  
+            subject_image = subject_image.to(device=device, dtype=dtype)
+            subject_image_latents = self._encode_vae_image(image=subject_image, generator=generator)
+            subject_latents = self._pack_latents(subject_image_latents, batch_size, num_channels_latents, height_cond*sub_number, width_cond)
+            mask2 = torch.zeros(shape_subject, device=device, dtype=dtype)
+            mask2 = self._pack_latents(mask2, batch_size, num_channels_latents, height_cond*sub_number, width_cond)
+            latent_subject_ids = prepare_latent_subject_ids(height_cond, width_cond, device, dtype)
+            latent_subject_ids[:, 1] += 64  # fixed offset
+            subject_latent_image_ids = torch.concat([latent_subject_ids for _ in range(sub_number)], dim=-2)
+            latents_to_concat.append(subject_latents)
+            latents_ids_to_concat.append(subject_latent_image_ids)
+            
+        # spatial
+        if condition_image is not None:
+            shape_cond = (batch_size, num_channels_latents, height_cond*cond_number, width_cond)  
+            condition_image = condition_image.to(device=device, dtype=dtype)
+            self.vae = self.vae.to(dtype)
+            image_latents = self._encode_vae_image(image=condition_image, generator=generator)
+            cond_latents = self._pack_latents(image_latents, batch_size, num_channels_latents, height_cond*cond_number, width_cond)
+            mask3 = torch.zeros(shape_cond, device=device, dtype=dtype)
+            mask3 = self._pack_latents(mask3, batch_size, num_channels_latents, height_cond*cond_number, width_cond) 
+            cond_latent_image_ids = cond_latent_image_ids
+            cond_latent_image_ids = torch.concat([cond_latent_image_ids for _ in range(cond_number)], dim=-2)
+            latents_ids_to_concat.append(cond_latent_image_ids)
+            latents_to_concat.append(cond_latents)
+
+        cond_latents = torch.concat(latents_to_concat, dim=-2)
+        latent_image_ids = torch.concat(latents_ids_to_concat, dim=-2)
+        return cond_latents, latent_image_ids, noise_latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def __call__(
+            self,
+            args: Any = None,  
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            num_inference_steps: int = 28,
+            timesteps: List[int] = None,
+            guidance_scale: float = 3.5,
+            num_images_per_prompt: Optional[int] = 1,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            output_type: Optional[str] = "pil",
+            return_dict: bool = True,
+            joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+            callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+            callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+            max_sequence_length: int = 512,
+            spatial_images=None,
+            subject_images=None,
+            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
+        
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+        
+        cond_number = len(spatial_images)
+        sub_number = len(subject_images)
+        
+        if sub_number > 0:
+            subject_image_ls = []
+            for subject_image in subject_images:
+                w, h = subject_image.size[:2]
+                scale = self.cond_size / max(h, w)
+                new_h, new_w = int(h * scale), int(w * scale)
+                subject_image = self.image_processor.preprocess(subject_image, height=new_h, width=new_w)
+                subject_image = subject_image.to(dtype=torch.float32)
+                pad_h = cond_size - subject_image.shape[-2]
+                pad_w = cond_size - subject_image.shape[-1]
+                subject_image = pad(
+                    subject_image,
+                    padding=(int(pad_w / 2), int(pad_h / 2), int(pad_w / 2), int(pad_h / 2)),
+                    fill=0
+                )
+                subject_image_ls.append(subject_image)
+            subject_image = torch.concat(subject_image_ls, dim=-2)
+        else:
+            subject_image = None
+        
+        if cond_number > 0:
+            condition_image_ls = []
+            for img in spatial_images:
+                condition_image = self.image_processor.preprocess(img, height=self.cond_size, width=self.cond_size)
+                condition_image = condition_image.to(dtype=torch.float32)
+                condition_image_ls.append(condition_image)
+            condition_image = torch.concat(condition_image_ls, dim=-2)
+        else:
+            condition_image = None
+        
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4  # 16
+        cond_latents, latent_image_ids, noise_latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            subject_image,
+            condition_image,
+            latents,
+            cond_number,
+            sub_number
+        )
+        latents = noise_latents
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                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(
+                    hidden_states=latents,  # 1 4096 64
+                    cond_hidden_states=cond_latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    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
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                latents = latents
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents.to(dtype=self.vae.dtype), return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

train/src/prompt_helper.py

@@ -0,0 +1,215 @@
+import torch
+import os 
+import os.path as osp 
+
+
+def load_text_encoders(args, class_one, class_two):
+    text_encoder_one = class_one.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant
+    )
+    text_encoder_two = class_two.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder_2", revision=args.revision, variant=args.variant
+    )
+    return text_encoder_one, text_encoder_two
+
+
+def tokenize_prompt(tokenizer, prompt, max_sequence_length):
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    return text_input_ids
+
+
+def tokenize_prompt_clip(tokenizer, prompt):
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=77,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    return text_input_ids
+
+
+def tokenize_prompt_t5(tokenizer, prompt):
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=512,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    return text_input_ids
+
+
+def _encode_prompt_with_t5(
+        text_encoder,
+        tokenizer,
+        max_sequence_length=512,
+        prompt=None,
+        num_images_per_prompt=1,
+        device=None,
+        text_input_ids=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device))[0]
+
+    dtype = text_encoder.dtype
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    _, seq_len, _ = prompt_embeds.shape
+
+    # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+    return prompt_embeds
+
+
+def _encode_prompt_with_clip(
+        text_encoder,
+        tokenizer,
+        prompt: str,
+        device=None,
+        text_input_ids=None,
+        num_images_per_prompt: int = 1,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+    # Use pooled output of CLIPTextModel
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
+
+    # duplicate text embeddings for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+    return prompt_embeds
+
+
+def encode_prompt(
+        args, 
+        text_encoders,
+        tokenizers,
+        prompt: str,
+        max_sequence_length,
+        device=None,
+        num_images_per_prompt: int = 1,
+        text_input_ids_list=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    dtype = text_encoders[0].dtype
+
+    _prompt_ = "_".join(prompt) 
+    if osp.exists(osp.join(args.inference_embeds_dir, f"{_prompt_}.pth")): 
+        prompt_embeds = torch.load(osp.join(args.inference_embeds_dir, f"{_prompt_}.pth")) 
+        pooled_prompt_embeds = prompt_embeds["pooled_prompt_embeds"] 
+        prompt_embeds = prompt_embeds["prompt_embeds"] 
+
+    else: 
+        pooled_prompt_embeds = _encode_prompt_with_clip(
+            text_encoder=text_encoders[0],
+            tokenizer=tokenizers[0],
+            prompt=prompt,
+            device=device if device is not None else text_encoders[0].device,
+            num_images_per_prompt=num_images_per_prompt,
+            text_input_ids=text_input_ids_list[0] if text_input_ids_list else None,
+        )
+
+        prompt_embeds = _encode_prompt_with_t5(
+            text_encoder=text_encoders[1],
+            tokenizer=tokenizers[1],
+            max_sequence_length=max_sequence_length,
+            prompt=prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device if device is not None else text_encoders[1].device,
+            text_input_ids=text_input_ids_list[1] if text_input_ids_list else None,
+        )
+
+    text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+
+
+def encode_token_ids(text_encoders, tokens, accelerator, num_images_per_prompt=1, device=None):
+    text_encoder_clip = text_encoders[0]
+    text_encoder_t5 = text_encoders[1]
+    tokens_clip, tokens_t5 = tokens[0], tokens[1]
+    batch_size = tokens_clip.shape[0]
+
+    if device == "cpu":
+        device = "cpu"
+    else:
+        device = accelerator.device
+
+    # clip
+    prompt_embeds = text_encoder_clip(tokens_clip.to(device), output_hidden_states=False)
+    # Use pooled output of CLIPTextModel
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=text_encoder_clip.dtype, device=accelerator.device)
+    # duplicate text embeddings for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    pooled_prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+    pooled_prompt_embeds = pooled_prompt_embeds.to(dtype=text_encoder_clip.dtype, device=accelerator.device)
+
+    # t5
+    prompt_embeds = text_encoder_t5(tokens_t5.to(device))[0]
+    dtype = text_encoder_t5.dtype
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=accelerator.device)
+    _, seq_len, _ = prompt_embeds.shape
+    # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+    text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=accelerator.device, dtype=dtype)
+
+    return prompt_embeds, pooled_prompt_embeds, text_ids

train/src/transformer_flux.py

@@ -0,0 +1,603 @@
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import os 
+import os.path as osp 
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    FluxAttnProcessor2_0,
+    FluxAttnProcessor2_0_NPU,
+    FusedFluxAttnProcessor2_0,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+@maybe_allow_in_graph
+class FluxSingleTransformerBlock(nn.Module):
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        if is_torch_npu_available():
+            processor = FluxAttnProcessor2_0_NPU()
+        else:
+            processor = FluxAttnProcessor2_0()
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=processor,
+            qk_norm="rms_norm",
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        cond_temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> torch.Tensor:
+        use_cond = cond_hidden_states is not None
+        
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+        
+        if use_cond:
+            residual_cond = cond_hidden_states
+            norm_cond_hidden_states, cond_gate = self.norm(cond_hidden_states, emb=cond_temb)
+            mlp_cond_hidden_states = self.act_mlp(self.proj_mlp(norm_cond_hidden_states))
+        
+        norm_hidden_states_concat = torch.concat([norm_hidden_states, norm_cond_hidden_states], dim=-2)
+        
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states_concat,
+            image_rotary_emb=image_rotary_emb,
+            use_cond=use_cond,
+            **joint_attention_kwargs,
+        )
+        if use_cond:
+            attn_output, cond_attn_output = attn_output
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        if use_cond:
+            condition_latents = torch.cat([cond_attn_output, mlp_cond_hidden_states], dim=2)
+            cond_gate = cond_gate.unsqueeze(1)
+            condition_latents = cond_gate * self.proj_out(condition_latents)
+            condition_latents = residual_cond + condition_latents
+        
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return hidden_states, condition_latents if use_cond else None
+
+
+@maybe_allow_in_graph
+class FluxTransformerBlock(nn.Module):
+    def __init__(
+        self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+    ):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = FluxAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        cond_temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        use_cond = cond_hidden_states is not None
+        
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        if use_cond:
+                (
+                    norm_cond_hidden_states,
+                    cond_gate_msa,
+                    cond_shift_mlp,
+                    cond_scale_mlp,
+                    cond_gate_mlp,
+                ) = self.norm1(cond_hidden_states, emb=cond_temb)
+                
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+        
+        norm_hidden_states = torch.concat([norm_hidden_states, norm_cond_hidden_states], dim=-2)
+        
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        # Attention.
+        attention_outputs = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            use_cond=use_cond,
+            **joint_attention_kwargs,
+        )
+
+        attn_output, context_attn_output = attention_outputs[:2]
+        cond_attn_output = attention_outputs[2] if use_cond else None
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        if use_cond:
+            cond_attn_output = cond_gate_msa.unsqueeze(1) * cond_attn_output
+            cond_hidden_states = cond_hidden_states + cond_attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        
+        if use_cond:
+            norm_cond_hidden_states = self.norm2(cond_hidden_states)
+            norm_cond_hidden_states = (
+                norm_cond_hidden_states * (1 + cond_scale_mlp[:, None])
+                + cond_shift_mlp[:, None]
+            )
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+        hidden_states = hidden_states + ff_output
+        
+        if use_cond:
+            cond_ff_output = self.ff(norm_cond_hidden_states)
+            cond_ff_output = cond_gate_mlp.unsqueeze(1) * cond_ff_output
+            cond_hidden_states = cond_hidden_states + cond_ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states, cond_hidden_states if use_cond else None
+
+
+class FluxTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, FluxTransformer2DLoadersMixin
+):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        out_channels: Optional[int] = None,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False,
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+    ):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedFluxAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_hidden_states: torch.Tensor = None,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+        call_ids: list = None,  
+        cuboids_segmasks: torch.Tensor = None, 
+        store_qk: bool = False, 
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        if cond_hidden_states is not None:
+            use_condition = True
+        else:
+            use_condition = False
+        
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+            joint_attention_kwargs = {}
+        joint_attention_kwargs["call_ids"] = call_ids 
+        joint_attention_kwargs["cuboids_segmasks"] = cuboids_segmasks 
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        hidden_states = self.x_embedder(hidden_states)
+        cond_hidden_states = self.x_embedder(cond_hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        
+        cond_temb = (
+            self.time_text_embed(torch.ones_like(timestep) * 0, pooled_projections)
+                if guidance is None
+                else self.time_text_embed(
+                    torch.ones_like(timestep) * 0, guidance, pooled_projections
+                )
+            )
+        
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    cond_temb=cond_temb if use_condition else None,
+                    cond_hidden_states=cond_hidden_states if use_condition else None,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                if store_qk: 
+                    overall_block_idx = index_block 
+                    joint_attention_kwargs["store_qk"] = osp.join(store_qk, f"{str(overall_block_idx).zfill(3)}") 
+
+                encoder_hidden_states, hidden_states, cond_hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cond_hidden_states=cond_hidden_states if use_condition else None,
+                    temb=temb,
+                    cond_temb=cond_temb if use_condition else None,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        # note that the encoder_hidden_states are concatenated in FRONT of the hidden states, not BEHIND 
+        # this would change the attention mask calculation.   
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states, cond_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    cond_temb=cond_temb if use_condition else None,
+                    cond_hidden_states=cond_hidden_states if use_condition else None,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                if store_qk: 
+                    overall_block_idx = index_block + len(self.transformer_blocks) 
+                    joint_attention_kwargs["store_qk"] = osp.join(store_qk, f"{str(overall_block_idx).zfill(3)}") 
+
+                hidden_states, cond_hidden_states = block(
+                    hidden_states=hidden_states,
+                    cond_hidden_states=cond_hidden_states if use_condition else None,
+                    temb=temb,
+                    cond_temb=cond_temb if use_condition else None,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

train/train.py

@@ -0,0 +1,1463 @@
+import argparse
+import copy
+import logging
+import random 
+import math
+import os
+import shutil
+import gc 
+from contextlib import nullcontext
+from pathlib import Path
+import re
+from safetensors.torch import save_file
+
+from PIL import Image
+import numpy as np
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
+
+from tqdm.auto import tqdm
+from transformers import CLIPTokenizer, PretrainedConfig, T5TokenizerFast
+
+import diffusers
+
+from diffusers import (
+	AutoencoderKL,
+	FlowMatchEulerDiscreteScheduler
+)
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import (
+	cast_training_params,
+	compute_density_for_timestep_sampling,
+	compute_loss_weighting_for_sd3,
+)
+import os.path as osp 
+from diffusers.utils.torch_utils import is_compiled_module
+from diffusers.utils import (
+	check_min_version,
+	is_wandb_available,
+	convert_unet_state_dict_to_peft
+)
+
+from src.prompt_helper import *
+from src.lora_helper import *
+from src.pipeline import FluxPipeline, resize_position_encoding, prepare_latent_subject_ids
+from src.layers import MultiDoubleStreamBlockLoraProcessor, MultiSingleStreamBlockLoraProcessor
+from src.transformer_flux import FluxTransformer2DModel
+from src.jsonl_datasets import make_train_dataset, collate_fn
+
+if is_wandb_available():
+	import wandb
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0.dev0")
+
+logger = get_logger(__name__)
+
+import matplotlib.pyplot as plt
+import torch
+
+def create_validation_figure(output_image, spatial_image, subject_image, prompt, validation_idx, global_step):
+	"""
+	Create a 2x2 matplotlib figure showing validation results.
+	
+	Args:
+		output_image: Generated output image (PIL Image)
+		spatial_image: Spatial condition image (PIL Image or None)
+		subject_image: Subject condition image (PIL Image or None) 
+		prompt: Text prompt string
+		validation_idx: Index of validation prompt
+		global_step: Current training step
+		
+	Returns:
+		matplotlib figure
+	"""
+	fig, axes = plt.subplots(2, 2, figsize=(12, 20))
+	fig.suptitle(f'Validation Results - Step {global_step} - Prompt {validation_idx}', fontsize=14)
+	
+	# Output image (top-left)
+	axes[0, 0].imshow(np.array(output_image))
+	axes[0, 0].set_title('Generated Output')
+	axes[0, 0].axis('off')
+	
+	# Spatial condition (top-right)
+	if spatial_image is not None:
+		axes[0, 1].imshow(np.array(spatial_image))
+		axes[0, 1].set_title('Spatial Condition')
+	else:
+		axes[0, 1].text(0.5, 0.5, 'NOT AVAILABLE', 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[0, 1].transAxes, fontsize=14, fontweight='bold')
+		axes[0, 1].set_title('Spatial Condition')
+	axes[0, 1].axis('off')
+	
+	# Subject condition (bottom-left)
+	if subject_image is not None:
+		axes[1, 0].imshow(np.array(subject_image))
+		axes[1, 0].set_title('Subject Condition')
+	else:
+		axes[1, 0].text(0.5, 0.5, 'NOT AVAILABLE', 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[1, 0].transAxes, fontsize=14, fontweight='bold')
+		axes[1, 0].set_title('Subject Condition')
+	axes[1, 0].axis('off')
+	
+	# Prompt and info (bottom-right)
+	info_text = f'Prompt:\n"{prompt}"\n\nStep: {global_step}\nValidation Index: {validation_idx}'
+	axes[1, 1].text(0.5, 0.5, info_text,
+				   horizontalalignment='center', verticalalignment='center',
+				   transform=axes[1, 1].transAxes, fontsize=10, wrap=True)
+	axes[1, 1].set_title('Prompt & Info')
+	axes[1, 1].axis('off')
+	
+	plt.tight_layout()
+	return fig
+
+
+def visualize_training_data(batch, vae, model_input, noisy_model_input, cond_input, subject_input, args, global_step, accelerator):
+	"""
+	Visualize training data including all entities from the batch.
+	
+	Args:
+		batch: Training batch containing data
+		vae: VAE model for decoding latents
+		model_input: Clean latents before adding noise
+		noisy_model_input: Noisy latents passed to transformer
+		cond_input: Spatial condition latents (may be None)
+		subject_input: Subject condition latents (may be None)
+		args: Training arguments
+		global_step: Current training step
+		accelerator: Accelerator instance
+	"""
+	
+	# Check availability of conditions
+	has_spatial_condition = batch["cond_pixel_values"] is not None
+	has_subject_condition = batch["subject_pixel_values"] is not None
+	has_cuboids_segmasks = "cuboids_segmasks" in batch and batch["cuboids_segmasks"] is not None
+	has_cuboids_segmasks_bev = "cuboids_segmasks_bev" in batch and batch["cuboids_segmasks_bev"] is not None
+	
+	# Initialize variables
+	spatial_img = None
+	subject_img = None
+	
+	with torch.no_grad():
+		# Get VAE config for proper decoding
+		vae_config_shift_factor = vae.config.shift_factor
+		vae_config_scaling_factor = vae.config.scaling_factor
+		vae_dtype = vae.dtype 
+		vae = vae.to(torch.float32) 
+		
+		# Decode spatial condition if available
+		if has_spatial_condition:
+			cond_for_decode = (cond_input / vae_config_scaling_factor) + vae_config_shift_factor
+			spatial_decoded = vae.decode(cond_for_decode.float()).sample
+			spatial_decoded = (spatial_decoded / 2 + 0.5).clamp(0, 1)  # Normalize to [0,1]
+			spatial_img = spatial_decoded[0].float().cpu().permute(1, 2, 0).numpy()
+		
+		# Decode subject condition if available
+		if has_subject_condition:
+			subject_for_decode = (subject_input / vae_config_scaling_factor) + vae_config_shift_factor
+			subject_decoded = vae.decode(subject_for_decode.float()).sample
+			subject_decoded = (subject_decoded / 2 + 0.5).clamp(0, 1)  # Normalize to [0,1]
+			subject_img = subject_decoded[0].float().cpu().permute(1, 2, 0).numpy()
+		
+		# Decode clean model input
+		clean_for_decode = (model_input / vae_config_scaling_factor) + vae_config_shift_factor
+		clean_decoded = vae.decode(clean_for_decode.float()).sample
+		clean_decoded = (clean_decoded / 2 + 0.5).clamp(0, 1)
+		
+		# Decode noisy model input
+		noisy_for_decode = (noisy_model_input / vae_config_scaling_factor) + vae_config_shift_factor
+		noisy_decoded = vae.decode(noisy_for_decode.float()).sample
+		noisy_decoded = (noisy_decoded / 2 + 0.5).clamp(0, 1)
+		
+		# Convert to CPU and numpy for visualization (take first batch item)
+		clean_img = clean_decoded[0].float().cpu().permute(1, 2, 0).numpy()
+		noisy_img = noisy_decoded[0].float().cpu().permute(1, 2, 0).numpy()
+		
+		# Get text prompt and other info
+		text_prompt = batch["prompts"][0] if isinstance(batch["prompts"], list) else batch["prompts"]
+		call_id = batch["call_ids"][0] if batch["call_ids"] is not None else "N/A"
+		
+		# Create figure with more subplots to accommodate all entities including BEV
+		fig, axes = plt.subplots(4, 3, figsize=(18, 24))
+		# fig.suptitle(f'Training Data Visualization - Step {global_step}', fontsize=16)
+		
+		# Spatial condition (0,0)
+		if has_spatial_condition and spatial_img is not None:
+			axes[0, 0].imshow(spatial_img)
+			axes[0, 0].set_title('Spatial Condition')
+		else:
+			axes[0, 0].text(0.5, 0.5, 'NOT AVAILABLE', 
+						   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[0, 0].transAxes, fontsize=14, fontweight='bold')
+			axes[0, 0].set_title('Spatial Condition')
+		axes[0, 0].axis('off')
+		
+		# Subject condition (0,1)
+		if has_subject_condition and subject_img is not None:
+			axes[0, 1].imshow(subject_img)
+			axes[0, 1].set_title('Subject Condition')
+		else:
+			axes[0, 1].text(0.5, 0.5, 'NOT AVAILABLE', 
+						   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[0, 1].transAxes, fontsize=14, fontweight='bold')
+			axes[0, 1].set_title('Subject Condition')
+		axes[0, 1].axis('off')
+		
+		# Clean model input (0,2)
+		axes[0, 2].imshow(clean_img)
+		axes[0, 2].set_title('Clean Model Input')
+		axes[0, 2].axis('off')
+		
+		# Noisy model input (1,0)
+		axes[1, 0].imshow(noisy_img)
+		axes[1, 0].set_title('Noisy Model Input')
+		axes[1, 0].axis('off')
+		
+		# Cuboids segmentation masks with legend (1,1 and 1,2)
+		if has_cuboids_segmasks:
+			segmask = batch["cuboids_segmasks"][0].float().cpu().numpy()  # Shape: (n_subjects, h, w)
+			n_subjects, h, w = segmask.shape
+			
+			# Only use first 4 subjects for visualization
+			n_subjects_to_show = min(4, n_subjects)
+			
+			# Create colored segmentation visualization
+			np.random.seed(42)  # For consistent colors
+			colors = np.random.rand(n_subjects_to_show + 1, 3)  # +1 for background
+			colors[0] = [0, 0, 0]  # Background is black
+			
+			# Create 2x2 grid of individual subject masks
+			grid_h, grid_w = 2, 2
+			combined_mask = np.zeros((h * grid_h, w * grid_w, 3))
+			
+			for idx in range(n_subjects_to_show):
+				row = idx // grid_w
+				col = idx % grid_w
+				
+				# Create binary mask for this subject
+				subject_mask = np.zeros((h, w, 3))
+				mask = segmask[idx] > 0.5  # Binary threshold
+				subject_mask[mask] = colors[idx + 1]
+				
+				# Place in grid
+				combined_mask[row*h:(row+1)*h, col*w:(col+1)*w] = subject_mask
+			
+			axes[1, 1].imshow(combined_mask)
+			axes[1, 1].set_title('Cuboids Segmentation (2x2 Grid)')
+			axes[1, 1].axis('off')
+			
+			# Create legend in the next subplot (1,2) - only for first 4 subjects
+			axes[1, 2].set_xlim(0, 1)
+			axes[1, 2].set_ylim(0, 1)
+			
+			# Add legend entries
+			legend_y_positions = np.linspace(0.9, 0.1, n_subjects_to_show + 1)
+			axes[1, 2].text(0.1, legend_y_positions[0], f"Background", 
+						   color=colors[0], fontsize=12, fontweight='bold')
+			
+			for subject_idx in range(n_subjects_to_show):
+				axes[1, 2].text(0.1, legend_y_positions[subject_idx + 1], 
+							   f"Subject {subject_idx}", 
+							   color=colors[subject_idx + 1], fontsize=12, fontweight='bold')
+			
+			axes[1, 2].set_title('Segmentation Legend (First 4)')
+			axes[1, 2].axis('off')
+		else:
+			axes[1, 1].text(0.5, 0.5, 'NOT AVAILABLE', 
+						   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[1, 1].transAxes, fontsize=14, fontweight='bold')
+			axes[1, 1].set_title('Cuboids Segmentation')
+			axes[1, 1].axis('off')
+			
+			axes[1, 2].text(0.5, 0.5, 'NOT AVAILABLE', 
+						   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[1, 2].transAxes, fontsize=14, fontweight='bold')
+			axes[1, 2].set_title('Segmentation Legend')
+			axes[1, 2].axis('off')
+		
+		# BEV Cuboids segmentation masks with legend (2,0 and 2,1)
+		if has_cuboids_segmasks_bev:
+			segmask_bev = batch["cuboids_segmasks_bev"][0].float().cpu().numpy()  # Shape: (n_subjects, h, w)
+			n_subjects_bev, h_bev, w_bev = segmask_bev.shape
+			
+			# Create colored segmentation visualization for BEV (use different seed for different colors)
+			np.random.seed(123)  # Different seed for BEV colors
+			colors_bev = np.random.rand(n_subjects_bev + 1, 3)  # +1 for background
+			colors_bev[0] = [0, 0, 0]  # Background is black
+			
+			# Create RGB image from BEV segmentation
+			colored_segmask_bev = np.zeros((h_bev, w_bev, 3))
+			for subject_idx in range(n_subjects_bev):
+				mask_bev = segmask_bev[subject_idx] > 0.5  # Binary threshold
+				colored_segmask_bev[mask_bev] = colors_bev[subject_idx + 1]
+			
+			axes[2, 0].imshow(colored_segmask_bev)
+			axes[2, 0].set_title('BEV Cuboids Segmentation')
+			axes[2, 0].axis('off')
+			
+			# Create BEV legend in the next subplot (2,1)
+			axes[2, 1].set_xlim(0, 1)
+			axes[2, 1].set_ylim(0, 1)
+			
+			# Add BEV legend entries
+			legend_y_positions_bev = np.linspace(0.9, 0.1, n_subjects_bev + 1)
+			axes[2, 1].text(0.1, legend_y_positions_bev[0], f"Background", 
+						   color=colors_bev[0], fontsize=12, fontweight='bold')
+			
+			for subject_idx in range(n_subjects_bev):
+				axes[2, 1].text(0.1, legend_y_positions_bev[subject_idx + 1], 
+							   f"Subject {subject_idx}", 
+							   color=colors_bev[subject_idx + 1], fontsize=12, fontweight='bold')
+			
+			axes[2, 1].set_title('BEV Segmentation Legend')
+			axes[2, 1].axis('off')
+		else:
+			axes[2, 0].text(0.5, 0.5, 'NOT AVAILABLE', 
+					   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[2, 0].transAxes, fontsize=14, fontweight='bold')
+			axes[2, 0].set_title('BEV Cuboids Segmentation')
+			axes[2, 0].axis('off')
+		
+			axes[2, 1].text(0.5, 0.5, 'NOT AVAILABLE', 
+						   horizontalalignment='center', verticalalignment='center',
+						   transform=axes[2, 1].transAxes, fontsize=14, fontweight='bold')
+			axes[2, 1].set_title('BEV Segmentation Legend')
+			axes[2, 1].axis('off')
+		
+		# Text prompt and call ID (2,2)
+		axes[2, 2].text(0.5, 0.5, f'Text Prompt:\n\n"{text_prompt}"\n\nCall ID: {call_id}', 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[2, 2].transAxes, fontsize=12, wrap=True)
+		axes[2, 2].set_title('Text Prompt & Call ID')
+		axes[2, 2].axis('off')
+		
+		# Pixel values info (3,0)
+		pixel_info = f'Pixel Values Shape: {batch["pixel_values"].shape}\n'
+		if has_spatial_condition:
+			pixel_info += f'Spatial Shape: {batch["cond_pixel_values"].shape}\n'
+		if has_subject_condition:
+			pixel_info += f'Subject Shape: {batch["subject_pixel_values"].shape}\n'
+		if has_cuboids_segmasks:
+			pixel_info += f'Cuboids Segmasks: {len(batch["cuboids_segmasks"])}\n'
+		if has_cuboids_segmasks_bev:
+			pixel_info += f'BEV Segmasks: {len(batch["cuboids_segmasks_bev"])}'
+		
+		axes[3, 0].text(0.5, 0.5, pixel_info, 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[3, 0].transAxes, fontsize=10, fontfamily='monospace')
+		axes[3, 0].set_title('Tensor Shapes')
+		axes[3, 0].axis('off')
+		
+		# Training info (3,1)
+		training_info = f'Global Step: {global_step}\nConditions:\nSpatial: {"✓" if has_spatial_condition else "✗"}\nSubject: {"✓" if has_subject_condition else "✗"}\nSegmasks: {"✓" if has_cuboids_segmasks else "✗"}\nBEV Segmasks: {"✓" if has_cuboids_segmasks_bev else "✗"}'
+		axes[3, 1].text(0.5, 0.5, training_info, 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[3, 1].transAxes, fontsize=12, fontfamily='monospace')
+		axes[3, 1].set_title('Training Info')
+		axes[3, 1].axis('off')
+		
+		# Additional info (3,2) - can be used for any extra debugging info
+		axes[3, 2].text(0.5, 0.5, 'Additional Info\n(Reserved)', 
+					   horizontalalignment='center', verticalalignment='center',
+					   transform=axes[3, 2].transAxes, fontsize=12, fontfamily='monospace')
+		axes[3, 2].set_title('Reserved')
+		axes[3, 2].axis('off')
+		
+		plt.tight_layout()
+		
+		# Save the visualization
+		save_dir = os.path.join(args.output_dir, "visualizations")
+		os.makedirs(save_dir, exist_ok=True)
+		save_path = os.path.join(save_dir, f"training_vis_step_{global_step}.png")
+		plt.savefig(save_path, dpi=150, bbox_inches='tight')
+		plt.close()
+		
+		logger.info(f"Training visualization saved to {save_path}")
+
+		vae = vae.to(vae_dtype)
+
+def log_validation(
+		pipeline,
+		args,
+		accelerator,
+		pipeline_args,
+		step,
+		torch_dtype,
+		is_final_validation=False,
+):
+	logger.info(
+		f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+		f" {pipeline_args['prompt']}."
+	)
+	pipeline = pipeline.to(accelerator.device)
+	pipeline.set_progress_bar_config(disable=True)
+	# run inference
+	generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
+	# autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext()
+	autocast_ctx = nullcontext()
+
+	with autocast_ctx:
+		images = [pipeline(**pipeline_args, generator=generator).images[0] for _ in range(args.num_validation_images)]
+
+	# for tracker in accelerator.trackers:
+	# 	phase_name = "test" if is_final_validation else "validation"
+	# 	if tracker.name == "tensorboard":
+	# 		np_images = np.stack([np.asarray(img) for img in images])
+	# 		tracker.writer.add_images(phase_name, np_images, step, dataformats="NHWC")
+	# 	if tracker.name == "wandb":
+	# 		tracker.log(
+	# 			{
+	# 				phase_name: [
+	# 					wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images)
+	# 				]
+	# 			},
+	# 		)
+
+	return images
+
+
+def import_model_class_from_model_name_or_path(
+		pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder"
+):
+	text_encoder_config = PretrainedConfig.from_pretrained(
+		pretrained_model_name_or_path, subfolder=subfolder, revision=revision
+	)
+	model_class = text_encoder_config.architectures[0]
+	if model_class == "CLIPTextModel":
+		from transformers import CLIPTextModel
+
+		return CLIPTextModel
+	elif model_class == "T5EncoderModel":
+		from transformers import T5EncoderModel
+
+		return T5EncoderModel
+	else:
+		raise ValueError(f"{model_class} is not supported.")
+
+
+def parse_args(input_args=None):
+	parser = argparse.ArgumentParser(description="Simple example of a training script.")
+	parser.add_argument("--lora_num", type=int, default=2, help="number of the lora.")
+	parser.add_argument("--cond_size", type=int, default=512, help="size of the condition data.")
+	parser.add_argument("--test_h", type=int, default=1024, help="max side of the training data.")
+	parser.add_argument("--debug", type=int, default=0, help="whether to enter debug mode -- visualizations, gradient checks, etc.")
+	parser.add_argument("--test_w", type=int, default=1024, help="max side of the training data.")
+	parser.add_argument("--mode",type=str,default=None,help="The mode of the controller. Choose between ['depth', 'pose', 'canny'].")
+	parser.add_argument("--run_name",type=str,required=True,help="the name of the wandb run")
+	parser.add_argument(
+		"--train_data_dir",
+		type=str,
+		default="",
+		help=(
+			"A folder containing the training data. Folder contents must follow the structure described in"
+			" https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
+			" must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
+		),
+	)
+	parser.add_argument(
+		"--inference_embeds_dir",
+		type=str,
+		default="",
+		help=(
+			"the captions for images"
+		),
+	)
+	parser.add_argument(
+		"--pretrained_model_name_or_path",
+		type=str,
+		default="",
+		required=False,
+		help="Path to pretrained model or model identifier from huggingface.co/models.",
+	)
+	parser.add_argument(
+		"--pretrained_lora_path",
+		type=str,
+		default=None,
+		required=False,
+		help="Path to pretrained model",
+	)
+	parser.add_argument(
+		"--revision",
+		type=str,
+		default=None,
+		required=False,
+		help="Revision of pretrained model identifier from huggingface.co/models.",
+	)
+	parser.add_argument(
+		"--variant",
+		type=str,
+		default=None,
+		help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+	)
+	parser.add_argument(
+		"--spatial_column",
+		type=str,
+		default="None",
+		help="The column of the dataset containing the canny image. By "
+			 "default, the standard Image Dataset maps out 'file_name' "
+			 "to 'image'.",
+	)
+	parser.add_argument(
+		"--subject_column",
+		type=str,
+		default="image",
+		help="The column of the dataset containing the subject image. By "
+			 "default, the standard Image Dataset maps out 'file_name' "
+			 "to 'image'.",
+	)
+	parser.add_argument(
+		"--target_column",
+		type=str,
+		default="image",
+		help="The column of the dataset containing the target image. By "
+			 "default, the standard Image Dataset maps out 'file_name' "
+			 "to 'image'.",
+	)
+	parser.add_argument(
+		"--caption_column",
+		type=str,
+		default="caption_left,caption_right",
+		help="The column of the dataset containing the instance prompt for each image",
+	)
+	parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.")
+	parser.add_argument(
+		"--max_sequence_length",
+		type=int,
+		default=512,
+		help="Maximum sequence length to use with with the T5 text encoder",
+	)
+	parser.add_argument(
+		"--validation_prompt",
+		type=str,
+		nargs="+",
+		default="A woodenpot floating in a pool.",
+		help="A prompt that is used during validation to verify that the model is learning.",
+	)
+	parser.add_argument(
+		"--subject_test_images",
+		type=str,
+		nargs="+",
+		default=["/tiamat-NAS/zhangyuxuan/datasets/benchmark_dataset/decoritems_woodenpot/0.png"],
+		help="A list of subject test image paths.",
+	)
+	parser.add_argument(
+		"--spatial_test_images",
+		type=str,
+		nargs="+",
+		default=[],
+		help="A list of spatial test image paths.",
+	)
+	parser.add_argument(
+		"--num_validation_images",
+		type=int,
+		default=4,
+		help="Number of images that should be generated during validation with `validation_prompt`.",
+	)
+	parser.add_argument(
+		"--validation_steps",
+		type=int,
+		default=20,
+		help=(
+			"Run validation every X epochs. validation consists of running the prompt"
+			" `args.validation_prompt` multiple times: `args.num_validation_images`."
+		),
+	)
+	parser.add_argument(
+		"--ranks",
+		type=int,
+		nargs="+",
+		default=[128],
+		help=("The dimension of the LoRA update matrices."),
+	)
+	parser.add_argument(
+		"--network_alphas",
+		type=int,
+		nargs="+",
+		default=[128],
+		help=("The dimension of the LoRA update matrices."),
+	)
+	parser.add_argument(
+		"--output_dir",
+		type=str,
+		required=True,
+		help="The output directory where the model predictions and checkpoints will be written.",
+	)
+	parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+	parser.add_argument(
+		"--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader."
+	)
+	parser.add_argument("--num_train_epochs", type=int, default=50)
+	parser.add_argument(
+		"--max_train_steps",
+		type=int,
+		default=None,
+		help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+	)
+	parser.add_argument(
+		"--checkpointing_steps",
+		type=int,
+		default=1000,
+		help=(
+			"Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+			" checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+			" training using `--resume_from_checkpoint`."
+		),
+	)
+	parser.add_argument(
+		"--checkpoints_total_limit",
+		type=int,
+		default=None,
+		help=("Max number of checkpoints to store."),
+	)
+	parser.add_argument(
+		"--resume_from_checkpoint",
+		type=str,
+		default=None,
+		help=(
+			"Whether training should be resumed from a previous checkpoint. Use a path saved by"
+			' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+		),
+	)
+	parser.add_argument(
+		"--gradient_accumulation_steps",
+		type=int,
+		default=1,
+		help="Number of updates steps to accumulate before performing a backward/update pass.",
+	)
+	parser.add_argument(
+		"--gradient_checkpointing",
+		action="store_true",
+		help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+	)
+	parser.add_argument(
+		"--learning_rate",
+		type=float,
+		default=1e-4,
+		help="Initial learning rate (after the potential warmup period) to use.",
+	)
+
+	parser.add_argument(
+		"--guidance_scale",
+		type=float,
+		default=1,
+		help="the FLUX.1 dev variant is a guidance distilled model",
+	)
+	parser.add_argument(
+		"--scale_lr",
+		action="store_true",
+		default=False,
+		help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+	)
+	parser.add_argument(
+		"--lr_scheduler",
+		type=str,
+		default="constant",
+		help=(
+			'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+			' "constant", "constant_with_warmup"]'
+		),
+	)
+	parser.add_argument(
+		"--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+	)
+	parser.add_argument(
+		"--lr_num_cycles",
+		type=int,
+		default=1,
+		help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
+	)
+	parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
+	parser.add_argument(
+		"--dataloader_num_workers",
+		type=int,
+		default=2,
+		help=(
+			"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+		),
+	)
+	parser.add_argument(
+		"--weighting_scheme",
+		type=str,
+		default="none",
+		choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
+		help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
+	)
+	parser.add_argument(
+		"--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+	)
+	parser.add_argument(
+		"--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+	)
+	parser.add_argument(
+		"--mode_scale",
+		type=float,
+		default=1.29,
+		help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+	)
+	parser.add_argument(
+		"--optimizer",
+		type=str,
+		default="AdamW",
+		help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'),
+	)
+
+	parser.add_argument(
+		"--use_8bit_adam",
+		action="store_true",
+		help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW",
+	)
+
+	parser.add_argument(
+		"--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers."
+	)
+	parser.add_argument(
+		"--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers."
+	)
+	parser.add_argument(
+		"--prodigy_beta3",
+		type=float,
+		default=None,
+		help="coefficients for computing the Prodigy stepsize using running averages. If set to None, "
+			 "uses the value of square root of beta2. Ignored if optimizer is adamW",
+	)
+	parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+	parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params")
+	parser.add_argument(
+		"--adam_weight_decay_text_encoder", type=float, default=1e-03, help="Weight decay to use for text_encoder"
+	)
+
+	parser.add_argument(
+		"--adam_epsilon",
+		type=float,
+		default=1e-08,
+		help="Epsilon value for the Adam optimizer and Prodigy optimizers.",
+	)
+
+	parser.add_argument(
+		"--prodigy_use_bias_correction",
+		type=bool,
+		default=True,
+		help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW",
+	)
+	parser.add_argument(
+		"--prodigy_safeguard_warmup",
+		type=bool,
+		default=True,
+		help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. "
+			 "Ignored if optimizer is adamW",
+	)
+	parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+	parser.add_argument(
+		"--logging_dir",
+		type=str,
+		default="logs",
+		help=(
+			"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+			" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+		),
+	)
+	parser.add_argument(
+		"--cache_latents",
+		action="store_true",
+		default=False,
+		help="Cache the VAE latents",
+	)
+	parser.add_argument(
+		"--report_to",
+		type=str,
+		default="wandb",
+		help=(
+			'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+			' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+		),
+	)
+	parser.add_argument(
+		"--mixed_precision",
+		type=str,
+		default="bf16",
+		choices=["no", "fp16", "bf16"],
+		help=(
+			"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+			" 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+			" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+		),
+	)
+	parser.add_argument(
+		"--upcast_before_saving",
+		action="store_true",
+		default=False,
+		help=(
+			"Whether to upcast the trained transformer layers to float32 before saving (at the end of training). "
+			"Defaults to precision dtype used for training to save memory"
+		),
+	)
+
+	if input_args is not None:
+		args = parser.parse_args(input_args)
+	else:
+		args = parser.parse_args()
+	return args
+
+
+def main(args):
+	if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
+		# due to pytorch#99272, MPS does not yet support bfloat16.
+		raise ValueError(
+			"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+		)
+
+	if args.pretrained_lora_path is not None: 
+		assert osp.exists(args.pretrained_lora_path), f"Make sure that the `pretrained_lora_path` {args.pretrained_lora_path} exists." 
+		args.resume_from_checkpoint = osp.dirname(args.pretrained_lora_path) 
+
+	args.output_dir = osp.join(args.output_dir, args.run_name) 
+	args.logging_dir = osp.join(args.output_dir, args.logging_dir) 
+	os.makedirs(args.output_dir, exist_ok=True)
+	os.makedirs(args.logging_dir, exist_ok=True)
+	logging_dir = Path(args.output_dir, args.logging_dir)
+
+	if args.subject_column == "None":
+		args.subject_column = None
+	if args.spatial_column == "None":
+		args.spatial_column = None
+
+	accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+	# kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+	accelerator = Accelerator(
+		gradient_accumulation_steps=args.gradient_accumulation_steps,
+		mixed_precision=args.mixed_precision,
+		log_with=args.report_to,
+		project_config=accelerator_project_config,
+		# kwargs_handlers=[kwargs],
+	)
+
+	def save_model_hook(models, weights, output_dir):
+		pass 
+
+	def load_model_hook(models, input_dir):
+		pass 
+
+	# Disable AMP for MPS.
+	if torch.backends.mps.is_available():
+		accelerator.native_amp = False
+
+	if args.report_to == "wandb":
+		if not is_wandb_available():
+			raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+	# Make one log on every process with the configuration for debugging.
+	logging.basicConfig(
+		format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+		datefmt="%m/%d/%Y %H:%M:%S",
+		level=logging.INFO,
+	)
+	logger.info(accelerator.state, main_process_only=False)
+	if accelerator.is_local_main_process:
+		transformers.utils.logging.set_verbosity_warning()
+		diffusers.utils.logging.set_verbosity_info()
+	else:
+		transformers.utils.logging.set_verbosity_error()
+		diffusers.utils.logging.set_verbosity_error()
+
+	# If passed along, set the training seed now.
+	if args.seed is not None:
+		set_seed(args.seed)
+
+	# Handle the repository creation
+	if accelerator.is_main_process:
+		if args.output_dir is not None:
+			os.makedirs(args.output_dir, exist_ok=True)
+
+	# Load the tokenizers
+	tokenizer_one = CLIPTokenizer.from_pretrained(
+		args.pretrained_model_name_or_path,
+		subfolder="tokenizer",
+		revision=args.revision,
+	)
+	tokenizer_two = T5TokenizerFast.from_pretrained(
+		args.pretrained_model_name_or_path,
+		subfolder="tokenizer_2",
+		revision=args.revision,
+	)
+
+	# Load scheduler and models
+	noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+		args.pretrained_model_name_or_path, subfolder="scheduler"
+	)
+	noise_scheduler_copy = copy.deepcopy(noise_scheduler)
+	gc.collect() 
+	torch.cuda.empty_cache() 
+	vae = AutoencoderKL.from_pretrained(
+		args.pretrained_model_name_or_path,
+		subfolder="vae",
+		revision=args.revision,
+		variant=args.variant,
+	)
+	transformer = FluxTransformer2DModel.from_pretrained(
+		args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant
+	)
+
+	# We only train the additional adapter LoRA layers
+	transformer.requires_grad_(True)
+	vae.requires_grad_(False)
+
+	# For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
+	# as these weights are only used for inference, keeping weights in full precision is not required.
+	weight_dtype = torch.float32
+	if accelerator.mixed_precision == "fp16":
+		weight_dtype = torch.float16
+	elif accelerator.mixed_precision == "bf16":
+		weight_dtype = torch.bfloat16
+
+	if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
+		# due to pytorch#99272, MPS does not yet support bfloat16.
+		raise ValueError(
+			"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
+		)
+
+	vae.to(accelerator.device, dtype=weight_dtype)
+	transformer.to(accelerator.device, dtype=weight_dtype)
+
+	if args.gradient_checkpointing:
+		transformer.enable_gradient_checkpointing()
+
+	#### lora_layers ####
+	if args.pretrained_lora_path is not None:
+		lora_path = args.pretrained_lora_path
+		checkpoint = load_checkpoint(lora_path)
+		lora_attn_procs = {}
+		double_blocks_idx = list(range(19))
+		single_blocks_idx = list(range(38))
+		number = 1
+		for name, attn_processor in transformer.attn_processors.items():
+			match = re.search(r'\.(\d+)\.', name)
+			if match:
+				layer_index = int(match.group(1))
+			
+			if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
+				lora_state_dicts = {}
+				for key, value in checkpoint.items():
+					# Match based on the layer index in the key (assuming the key contains layer index)
+					if re.search(r'\.(\d+)\.', key):
+						checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+						if checkpoint_layer_index == layer_index and key.startswith("transformer_blocks"):
+							lora_state_dicts[key] = value
+				
+				print("setting LoRA Processor for", name)
+				lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
+					dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
+				)
+				
+				# Load the weights from the checkpoint dictionary into the corresponding layers
+				for n in range(number):
+					lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
+					lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
+					lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
+					lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
+					lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
+					lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
+					lora_attn_procs[name].proj_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.down.weight', None)
+					lora_attn_procs[name].proj_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.up.weight', None)
+				
+			elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
+				
+				lora_state_dicts = {}
+				for key, value in checkpoint.items():
+					# Match based on the layer index in the key (assuming the key contains layer index)
+					if re.search(r'\.(\d+)\.', key):
+						checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
+						if checkpoint_layer_index == layer_index and key.startswith("single_transformer_blocks"):
+							lora_state_dicts[key] = value
+				
+				print("setting LoRA Processor for", name)        
+				lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
+					dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
+				)
+				
+				# Load the weights from the checkpoint dictionary into the corresponding layers
+				for n in range(number):
+					lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
+					lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
+					lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
+					lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
+					lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
+					lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
+			else:
+				lora_attn_procs[name] = FluxAttnProcessor2_0()
+	else:
+		lora_attn_procs = {}
+		double_blocks_idx = list(range(19))
+		single_blocks_idx = list(range(38))
+		for name, attn_processor in transformer.attn_processors.items():
+			match = re.search(r'\.(\d+)\.', name)
+			if match:
+				layer_index = int(match.group(1))
+			if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
+				lora_state_dicts = {}
+				print("setting LoRA Processor for", name)
+				lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
+					dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
+				)
+
+			elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
+				print("setting LoRA Processor for", name)
+				lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
+					dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
+				)
+
+			else:
+				lora_attn_procs[name] = attn_processor        
+	######################
+	transformer.set_attn_processor(lora_attn_procs)
+	transformer.train()
+	for n, param in transformer.named_parameters():
+		if '_lora' not in n:
+			param.requires_grad = False
+	print(sum([p.numel() for p in transformer.parameters() if p.requires_grad]) / 1000000, 'M parameters')
+
+	def unwrap_model(model):
+		model = accelerator.unwrap_model(model)
+		model = model._orig_mod if is_compiled_module(model) else model
+		return model
+
+	# Potentially load in the weights and states from a previous save
+	if args.resume_from_checkpoint:
+		foldername = osp.basename(args.resume_from_checkpoint) 
+		first_epoch = epoch = int(foldername.split("-")[1].split("__")[0])  
+		initial_global_step = global_step = int(foldername.split("-")[-1]) 
+	else:
+		initial_global_step = 0
+		global_step = 0
+		first_epoch = 0
+
+	if args.scale_lr:
+		args.learning_rate = (
+				args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+		)
+
+	# Make sure the trainable params are in float32.
+	if args.mixed_precision == "fp16":
+		models = [transformer]
+		# only upcast trainable parameters (LoRA) into fp32
+		cast_training_params(models, dtype=torch.float32)
+
+	# Optimization parameters
+	params_to_optimize = [p for p in transformer.parameters() if p.requires_grad]
+	transformer_parameters_with_lr = {"params": params_to_optimize, "lr": args.learning_rate}
+	print(sum([p.numel() for p in transformer.parameters() if p.requires_grad]) / 1000000, 'parameters')
+
+	optimizer_class = torch.optim.AdamW
+	optimizer = optimizer_class(
+		[transformer_parameters_with_lr],
+		betas=(args.adam_beta1, args.adam_beta2),
+		weight_decay=args.adam_weight_decay,
+		eps=args.adam_epsilon,
+	)
+
+	tokenizers = [tokenizer_one, tokenizer_two]
+
+	# # Dataset and DataLoaders creation:
+	# train_dataset = make_train_dataset(args, tokenizers, accelerator)
+	# train_dataloader = torch.utils.data.DataLoader(
+	# 	train_dataset,
+	# 	batch_size=args.train_batch_size,
+	# 	shuffle=True,
+	# 	collate_fn=collate_fn,
+	# 	num_workers=args.dataloader_num_workers,
+	# )
+
+	# now, we will define a dataset for each epoch to make it easier to save the state   
+	shuffled_jsonls = os.listdir(osp.dirname(args.train_data_dir)) 
+	base_jsonl_name = osp.basename(args.train_data_dir).replace(".jsonl", "") 
+	shuffled_jsonls = sorted([_ for _ in shuffled_jsonls if _.endswith('.jsonl') and "shuffled" in _ and base_jsonl_name in _])   
+	shuffled_jsonls = [osp.join(osp.dirname(args.train_data_dir), _) for _ in shuffled_jsonls] 
+	print(f"{shuffled_jsonls = }")
+	# exit(0) 
+	assert len(shuffled_jsonls) > 0, f"Make sure that there are shuffled jsonl files in {osp.dirname(args.train_data_dir)}" 
+	train_dataloaders = [] 
+	for epoch in range(args.num_train_epochs): # prepare dataloader for each epoch, irrespective of the resume state  
+		shuffled_idx = epoch % len(shuffled_jsonls) 
+		train_data_file = shuffled_jsonls[shuffled_idx] 
+		assert osp.exists(train_data_file), f"Make sure that the train data jsonl file {train_data_file} exists." 
+		args.current_train_data_dir = train_data_file 
+		train_dataset = make_train_dataset(args, tokenizers, accelerator) 
+		train_dataloader = torch.utils.data.DataLoader(
+			train_dataset,
+			batch_size=args.train_batch_size,
+			shuffle=False, # yayy!! reproducible experiments!
+			collate_fn=collate_fn,
+			num_workers=args.dataloader_num_workers,
+		)
+		train_dataloaders.append(train_dataloader) 
+
+	vae_config_shift_factor = vae.config.shift_factor
+	vae_config_scaling_factor = vae.config.scaling_factor
+
+	# Scheduler and math around the number of training steps.
+	overrode_max_train_steps = False
+	num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)  
+	if args.max_train_steps is None:
+		args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+		overrode_max_train_steps = True
+
+	lr_scheduler = get_scheduler(
+		args.lr_scheduler,
+		optimizer=optimizer,
+		num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+		num_training_steps=args.max_train_steps * accelerator.num_processes,
+		num_cycles=args.lr_num_cycles,
+		power=args.lr_power,
+	)
+
+
+	accelerator.register_save_state_pre_hook(save_model_hook)
+	accelerator.register_load_state_pre_hook(load_model_hook)
+	optimizer, lr_scheduler = accelerator.prepare(
+		optimizer, lr_scheduler 
+	)
+
+	print(f"before preparation, {len(train_dataloaders[0]) = }") 
+
+	prepared_train_dataloaders = [] 
+	for train_dataloader in train_dataloaders: 
+		prepared_train_dataloaders.append(accelerator.prepare(train_dataloader)) 
+	train_dataloaders = prepared_train_dataloaders 
+
+	print(f"after preparation, {len(train_dataloaders[0]) = }") 
+
+	if args.pretrained_lora_path is not None: 
+		accelerator.load_state(osp.dirname(args.pretrained_lora_path))  
+
+	# Explicitly move optimizer states to accelerator.device
+	for state in optimizer.state.values():
+		for k, v in state.items():
+			if isinstance(v, torch.Tensor):
+				state[k] = v.to(accelerator.device)
+
+	transformer = accelerator.prepare(transformer) 
+
+	# We need to recalculate our total training steps as the size of the training dataloader may have changed.
+	num_update_steps_per_epoch = math.ceil(len(train_dataloaders[0]) / args.gradient_accumulation_steps)
+	if overrode_max_train_steps:
+		args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+	# Afterwards we recalculate our number of training epochs
+	args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+	# We need to initialize the trackers we use, and also store our configuration.
+	# The trackers initializes automatically on the main process.
+	# if accelerator.is_main_process:
+	# 	tracker_name = "Easy_Control"
+	# 	accelerator.init_trackers(tracker_name, config=vars(args))
+
+	if accelerator.is_main_process:  
+		tracker_config = vars(copy.deepcopy(args))
+		# tracker_config.pop("validation_images")
+		wandb_args = {
+			"wandb": {
+				"entity": "generative_parts", 
+				"name": args.run_name, 
+			}
+		} 
+		accelerator.init_trackers("seethrough3d", config=tracker_config, init_kwargs=wandb_args) 
+
+
+	# Train!
+	total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+	logger.info("***** Running training *****")
+	logger.info(f"  Num examples = {len(train_dataset)}")
+	logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+	logger.info(f"  Num Epochs = {args.num_train_epochs}")
+	logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+	logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+	logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+	logger.info(f"  Total optimization steps = {args.max_train_steps}")
+
+	progress_bar = tqdm(
+		range(0, args.max_train_steps),
+		initial=initial_global_step,
+		desc="Steps",
+		# Only show the progress bar once on each machine.
+		disable=not accelerator.is_local_main_process,
+	)
+
+	def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+		sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
+		schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
+		timesteps = timesteps.to(accelerator.device)
+		step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+		sigma = sigmas[step_indices].flatten()
+		while len(sigma.shape) < n_dim:
+			sigma = sigma.unsqueeze(-1)
+		return sigma
+	
+	# some fixed parameters 
+	vae_scale_factor = 16
+	height_cond = 2 * (args.cond_size // vae_scale_factor)
+	width_cond = 2 * (args.cond_size // vae_scale_factor)        
+	offset = 64
+
+	num_training_visualizations = 10   
+		
+	skip_steps = initial_global_step - first_epoch * num_update_steps_per_epoch   
+	print(f"{skip_steps = }")
+	for epoch in range(first_epoch, args.num_train_epochs):
+		transformer.train()
+		train_dataloader = train_dataloaders[epoch] # use a new dataloader for each epoch  
+		if epoch == first_epoch and skip_steps > 0:
+			logger.info(f"Skipping {skip_steps} batches in epoch {epoch} due to resuming from checkpoint")
+			# dataloader_iterator = skip_first_batches_manual(train_dataloader, skip_steps)
+			dataloader_iterator = accelerator.skip_first_batches(train_dataloader, skip_steps) 
+			# Convert back to enumerate format
+			enumerated_dataloader = enumerate(dataloader_iterator, start=skip_steps)
+		else: 
+			enumerated_dataloader = enumerate(train_dataloader) 
+		for step, batch in enumerated_dataloader: 
+			progress_bar.set_description(f"epoch {epoch}, dataset_ids: {batch['index']}") 
+			torch.cuda.empty_cache() 
+			models_to_accumulate = [transformer]
+			with accelerator.accumulate(models_to_accumulate):
+				
+				# tokens = [batch["text_ids_1"], batch["text_ids_2"]]
+				# prompt_embeds, pooled_prompt_embeds, text_ids = encode_token_ids(text_encoders, tokens, accelerator)
+				prompt_embeds = batch["prompt_embeds"] 
+				pooled_prompt_embeds = batch["pooled_prompt_embeds"] 
+				text_ids = torch.zeros((batch["prompt_embeds"].shape[1], 3)) 
+				prompt_embeds = prompt_embeds.to(dtype=vae.dtype, device=accelerator.device)
+				pooled_prompt_embeds = pooled_prompt_embeds.to(dtype=vae.dtype, device=accelerator.device)
+				text_ids = text_ids.to(dtype=vae.dtype, device=accelerator.device)
+				
+				pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
+				height_ = 2 * (int(pixel_values.shape[-2]) // vae_scale_factor)
+				width_ = 2 * (int(pixel_values.shape[-1]) // vae_scale_factor)
+
+				model_input = vae.encode(pixel_values).latent_dist.sample()
+				model_input = (model_input - vae_config_shift_factor) * vae_config_scaling_factor
+				model_input = model_input.to(dtype=weight_dtype)
+
+				latent_image_ids, cond_latent_image_ids = resize_position_encoding(
+					model_input.shape[0],
+					height_,
+					width_,
+					height_cond,
+					width_cond,
+					accelerator.device,
+					weight_dtype,
+				)
+
+				# Sample noise that we'll add to the latents
+				noise = torch.randn_like(model_input)
+				bsz = model_input.shape[0]
+
+				# Sample a random timestep for each image
+				# for weighting schemes where we sample timesteps non-uniformly
+				u = compute_density_for_timestep_sampling(
+					weighting_scheme=args.weighting_scheme,
+					batch_size=bsz,
+					logit_mean=args.logit_mean,
+					logit_std=args.logit_std,
+					mode_scale=args.mode_scale,
+				)
+				indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
+				timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)
+
+				# Add noise according to flow matching.
+				# zt = (1 - texp) * x + texp * z1
+				sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+				noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
+
+				packed_noisy_model_input = FluxPipeline._pack_latents(
+					noisy_model_input,
+					batch_size=model_input.shape[0],
+					num_channels_latents=model_input.shape[1],
+					height=model_input.shape[2],
+					width=model_input.shape[3],
+				)
+				
+				latent_image_ids_to_concat = [latent_image_ids]
+				packed_cond_model_input_to_concat = []
+				
+				if args.subject_column is not None:
+					# in case the condition is not spatial 
+					subject_pixel_values = batch["subject_pixel_values"].to(dtype=vae.dtype)
+					subject_input = vae.encode(subject_pixel_values).latent_dist.sample()
+					subject_input = (subject_input - vae_config_shift_factor) * vae_config_scaling_factor
+					subject_input = subject_input.to(dtype=weight_dtype)             
+					# the number of subjects in the concatenated subject image 
+					sub_number = subject_pixel_values.shape[-2] // args.cond_size
+					latent_subject_ids = prepare_latent_subject_ids(height_cond, width_cond, accelerator.device, weight_dtype)
+					latent_subject_ids[:, 1] += offset
+					sub_latent_image_ids = torch.concat([latent_subject_ids for _ in range(sub_number)], dim=-2)
+					latent_image_ids_to_concat.append(sub_latent_image_ids)
+					
+					packed_subject_model_input = FluxPipeline._pack_latents(    
+						subject_input,
+						batch_size=subject_input.shape[0],
+						num_channels_latents=subject_input.shape[1],
+						height=subject_input.shape[2],
+						width=subject_input.shape[3],
+					)
+					packed_cond_model_input_to_concat.append(packed_subject_model_input)
+				else: 
+					subject_input = None 
+				
+				if args.spatial_column is not None:
+					# in case the condition is spatial 
+					cond_pixel_values = batch["cond_pixel_values"].to(dtype=vae.dtype)             
+					cond_input = vae.encode(cond_pixel_values).latent_dist.sample()
+					cond_input = (cond_input - vae_config_shift_factor) * vae_config_scaling_factor
+					cond_input = cond_input.to(dtype=weight_dtype)
+					# number of conditions in the concatenated condition image 
+					cond_number = cond_pixel_values.shape[-2] // args.cond_size
+					cond_latent_image_ids = torch.concat([cond_latent_image_ids for _ in range(cond_number)], dim=-2)
+					latent_image_ids_to_concat.append(cond_latent_image_ids)
+
+					packed_cond_model_input = FluxPipeline._pack_latents(
+						cond_input,
+						batch_size=cond_input.shape[0],
+						num_channels_latents=cond_input.shape[1],
+						height=cond_input.shape[2],
+						width=cond_input.shape[3],
+					)
+					packed_cond_model_input_to_concat.append(packed_cond_model_input)
+				else: 
+					cond_input = None 
+					
+				latent_image_ids = torch.concat(latent_image_ids_to_concat, dim=-2)
+				cond_packed_noisy_model_input = torch.concat(packed_cond_model_input_to_concat, dim=-2)
+
+				# handle guidance
+				if accelerator.unwrap_model(transformer).config.guidance_embeds:
+					guidance = torch.tensor([args.guidance_scale], device=accelerator.device)
+					guidance = guidance.expand(model_input.shape[0])
+				else:
+					guidance = None
+
+				# Visualize training data before transformer forward pass
+				if accelerator.is_main_process and args.debug and num_training_visualizations > 0 and global_step % 5 == 0: 
+					visualize_training_data(
+						batch=batch,
+						vae=vae, 
+						model_input=model_input,
+						noisy_model_input=noisy_model_input,
+						cond_input=cond_input,
+						subject_input=subject_input,
+						args=args,
+						global_step=global_step,
+						accelerator=accelerator
+					)
+					num_training_visualizations -= 1 
+
+				# Predict the noise residual
+				model_pred = transformer(
+					hidden_states=packed_noisy_model_input,
+					# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+					cond_hidden_states=cond_packed_noisy_model_input,
+					timestep=timesteps / 1000,
+					guidance=guidance,
+					pooled_projections=pooled_prompt_embeds,
+					encoder_hidden_states=prompt_embeds,
+					txt_ids=text_ids,
+					img_ids=latent_image_ids,
+					return_dict=False,
+					call_ids=batch["call_ids"], 
+					cuboids_segmasks=batch["cuboids_segmasks"], 
+				)[0]
+
+				model_pred = FluxPipeline._unpack_latents(
+					model_pred,
+					height=int(pixel_values.shape[-2]),
+					width=int(pixel_values.shape[-1]),
+					vae_scale_factor=vae_scale_factor,
+				)
+
+				# these weighting schemes use a uniform timestep sampling
+				# and instead post-weight the loss
+				weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
+
+				# flow matching loss
+				target = noise - model_input
+
+				# Compute regular loss.
+				loss = torch.mean(
+					(weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
+					1,
+				)
+
+				loss = loss.mean()
+				accelerator.backward(loss)
+				if accelerator.sync_gradients:
+					params_to_clip = (transformer.parameters())
+					accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+
+				optimizer.step()
+				lr_scheduler.step()
+				optimizer.zero_grad()
+
+			# Checks if the accelerator has performed an optimization step behind the scenes
+			if accelerator.sync_gradients:
+				progress_bar.update(1)
+				global_step += 1
+
+				if accelerator.is_main_process:
+					if global_step % args.checkpointing_steps == 0:
+						# _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+						if args.checkpoints_total_limit is not None:
+							checkpoints = os.listdir(args.output_dir)
+							checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+							checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+							# before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+							if len(checkpoints) >= args.checkpoints_total_limit:
+								num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+								removing_checkpoints = checkpoints[0:num_to_remove]
+
+								logger.info(
+									f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+								)
+								logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+								for removing_checkpoint in removing_checkpoints:
+									removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+									shutil.rmtree(removing_checkpoint)
+
+						save_path = os.path.join(args.output_dir, f"epoch-{epoch}__checkpoint-{global_step}")
+						os.makedirs(save_path, exist_ok=True)
+						unwrapped_model_state = accelerator.unwrap_model(transformer).state_dict()
+						lora_state_dict = {k:unwrapped_model_state[k] for k in unwrapped_model_state.keys() if '_lora' in k}
+						save_file(
+							lora_state_dict,
+							os.path.join(save_path, "lora.safetensors")
+						)
+						accelerator.save_state(save_path) 
+						os.remove(osp.join(save_path, "model.safetensors")) 
+						logger.info(f"Saved state to {save_path}")
+
+			logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+			progress_bar.set_postfix(**logs)
+			accelerator.log(logs, step=global_step)
+
+		save_path = os.path.join(args.output_dir, f"epoch-{epoch}__checkpoint-{global_step}")
+		os.makedirs(save_path, exist_ok=True)
+		unwrapped_model_state = accelerator.unwrap_model(transformer).state_dict()
+		lora_state_dict = {k:unwrapped_model_state[k] for k in unwrapped_model_state.keys() if '_lora' in k}
+		save_file(
+			lora_state_dict,
+			os.path.join(save_path, "lora.safetensors")
+		)
+		accelerator.save_state(save_path) 
+		os.remove(osp.join(save_path, "model.safetensors")) 
+		logger.info(f"Saved state to {save_path}")
+		accelerator.wait_for_everyone()
+	accelerator.end_training()
+
+
+if __name__ == "__main__":
+	args = parse_args()
+	main(args)

train/train.sh

@@ -0,0 +1,71 @@
+#!/bin/bash
+#SBATCH --job-name=vaibhav
+#SBATCH --output=%j.out
+#SBATCH --ntasks=1                  
+#SBATCH --cpus-per-task=4           
+#SBATCH --mem=150G
+#SBATCH --gres=gpu:4                    
+#SBATCH --partition=ada
+
+# chetna
+# export MODEL_DIR="black-forest-labs/FLUX.1-Kontext-dev" # your flux path
+export MODEL_DIR="black-forest-labs/FLUX.1-dev" # your flux path
+export OUTPUT_DIR="/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids"  # your save path
+export CONFIG="./default_config.yaml"
+export TRAIN_DATA="/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv7_superhard/cuboids__upto_4subjects.jsonl" # your data jsonl file 
+export LOG_PATH="$OUTPUT_DIR/log"
+export INFERENCE_EMBEDS_DIR="/archive/vaibhav.agrawal/a-bev-of-the-latents/inference_embeds_datasetv7_superhard"
+
+export WANDB_API_KEY=f27c837d8d7d0c8d79f3eb1de21fa78233c03be6
+
+# kotak
+# export MODEL_DIR="black-forest-labs/FLUX.1-dev" # your flux path
+# export OUTPUT_DIR="/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids"  # your save path
+# export CONFIG="./default_config.yaml"
+# export TRAIN_DATA="/archive/vaibhav.agrawal/a-bev-of-the-latents/datasetv6/cuboids.jsonl" # your data jsonl file
+# export LOG_PATH="$OUTPUT_DIR/log"
+# export INFERENCE_EMBEDS_DIR="/archive/vaibhav.agrawal/a-bev-of-the-latents/inference_embeds_flux2"
+
+# kotak
+# export MODEL_DIR="black-forest-labs/FLUX.1-dev" # your flux path
+# export OUTPUT_DIR="./easycontrol_cuboids"  # your save path
+# export CONFIG="./default_config.yaml"
+# export TRAIN_DATA="/home/venky/vaibhav.agrawal/a-bev-of-the-latents/datasets/actual_data/datasetv6/cuboids.jsonl" # your data jsonl file
+# export LOG_PATH="$OUTPUT_DIR/log"
+# export INFERENCE_EMBEDS_DIR="/home/venky/vaibhav.agrawal/a-bev-of-the-latents/caching/inference_embeds_flux2"
+
+# i love this.
+accelerate launch --config_file $CONFIG train.py \
+    --pretrained_model_name_or_path $MODEL_DIR \
+    --cond_size=512 \
+    --subject_column="None" \
+    --spatial_column="cv" \
+    --target_column="target" \
+    --caption_column="caption" \
+    --ranks 128 \
+    --network_alphas 128 \
+    --lora_num 1 \
+    --output_dir=$OUTPUT_DIR \
+    --logging_dir=$LOG_PATH \
+    --run_name="rgb__r1" \
+    --debug=1 \
+    --mixed_precision="bf16" \
+    --train_data_dir=$TRAIN_DATA \
+    --learning_rate=1e-4 \
+    --train_batch_size=1 \
+    --inference_embeds_dir $INFERENCE_EMBEDS_DIR \
+    --validation_prompt "a photo of sedan and pickup truck and suv amongst autumn-colored trees along a winding river" "a photo of cow and suv on a sandy beach with palm trees swaying in the breeze" "a photo of table and horse and suv in a dense pine forest with tall trees reaching the sky" \
+    --num_train_epochs=1 \
+    --validation_steps=5000000000000 \
+    --checkpointing_steps=2500 \
+    --spatial_test_images "cuboids/sedan__pickup_truck__suv/005/cuboids.png" "cuboids/cow__suv/008/cuboids.png" "cuboids/table__horse__suv/007/cuboids.png" \
+    --subject_test_images None \
+    --test_h 512 \
+    --test_w 512 \
+    --num_validation_images=1
+
+    # --run_name="semantic_info_from_cuboid_cond" \
+    # --run_name="datasetv8__0.8_0.1_0.1" \
+    # --pretrained_lora_path="/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids/wireframe/epoch-0__checkpoint-5000/lora.safetensors" \
+    # --pretrained_lora_path="/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids/rgb/epoch-0__checkpoint-7500/lora.safetensors" \
+    # --pretrained_lora_path="/archive/vaibhav.agrawal/a-bev-of-the-latents/easycontrol_cuboids/datasetv9__wireframe_best_case/epoch-0__checkpoint-3888/lora.safetensors" \

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)