mem

ambient_diffusion.py

@@ -110,6 +110,7 @@ class AmbientDDPMPipeline(DDPMPipeline):
         guidance_scale: float = 1.5,
         output_type: str = "pt",
         return_dict: bool = True,
+        mem = False
     ):
         device = self.device
         h = w = self.unet.config.sample_size
@@ -126,6 +127,9 @@ class AmbientDDPMPipeline(DDPMPipeline):
 
         self.scheduler.set_timesteps(num_inference_steps, device=device)
 
+        if mem:
+            TCNP = torch.empty((batch_size, num_inference_steps), device=self.device)
+
         # Set up conditioning labels
         if class_labels is None:
             cond_lbls = torch.zeros(
@@ -138,7 +142,7 @@ class AmbientDDPMPipeline(DDPMPipeline):
 
         uncond_lbls = torch.zeros_like(cond_lbls)
 
-        for t in self.scheduler.timesteps:
+        for i, t in enumerate(self.scheduler.timesteps):
             # Always include mask channel to match training, but mask is all ones
             x_with_mask = torch.cat([x, dummy_mask], dim=1)
 
@@ -156,6 +160,10 @@ class AmbientDDPMPipeline(DDPMPipeline):
                 )
                 eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
 
+                if mem:
+                        # print(cond_output.squeeze().shape, uncond_output.shape)
+                        TCNP[:, i] = torch.linalg.norm(eps_cond.squeeze() - eps_uncond.squeeze(), dim=[0,1])
+
             # Standard update (same as before)
             x0_hat = eps
             sigma = self.scheduler._get_variance(t).sqrt()
@@ -167,4 +175,11 @@ class AmbientDDPMPipeline(DDPMPipeline):
         img = x.cpu().permute(0, 2, 3, 1).numpy()
         if output_type == "pil":
             img = self.numpy_to_pil(img)
-        return dict(images=img) if return_dict else (img,)
+            
+        if not return_dict:
+            return (img,)
+        
+        if mem:
+            return dict(images=img, TCNP=TCNP)
+
+        return dict(images=img)

conditional_ddpm_pipeline.py

@@ -15,6 +15,7 @@ class ConditionalDDPMPipeline(DDPMPipeline):
         class_labels=None,
         guidance_scale=1.5,
         return_dict=True,
+        mem=False,
     ):
         # Initialize with random noise (exactly like original)
         image = torch.randn(
@@ -31,8 +32,11 @@ class ConditionalDDPMPipeline(DDPMPipeline):
         # Setup the scheduler (exactly like original)
         self.scheduler.set_timesteps(num_inference_steps)
 
+        if mem:
+            TCNP = torch.empty((batch_size, num_inference_steps), device=self.device)
+
         # Denoising process
-        for t in self.scheduler.timesteps:
+        for i, t in enumerate(self.scheduler.timesteps):
             # Only difference is we pass class_labels to the model
             with torch.no_grad():
                 if guidance_scale > 1.0 and class_labels is not None:
@@ -51,6 +55,10 @@ class ConditionalDDPMPipeline(DDPMPipeline):
                     model_output = uncond_output + guidance_scale * (
                         cond_output - uncond_output
                     )
+
+                    if mem:
+                        # print(cond_output.squeeze().shape, uncond_output.shape)
+                        TCNP[:, i] = torch.linalg.norm(cond_output.squeeze() - uncond_output.squeeze(), dim=[0,1])
                 else:
                     # Standard pass with conditioning
                     model_output = self.unet(
@@ -71,5 +79,8 @@ class ConditionalDDPMPipeline(DDPMPipeline):
 
         if not return_dict:
             return (image,)
+        
+        if mem:
+            return dict(images=image, TCNP=TCNP)
 
         return dict(images=image)

handler.py

@@ -86,7 +86,7 @@ class EndpointHandler:
             scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")
             print("Created new default DDPMScheduler instance.")
 
-        if pipeline_type == "ambient":
+        if model_variant_name == "ambient":
             scheduler.config.prediction_type = "sample"
         else: 
             scheduler.config.prediction_type = "epsilon"
@@ -97,13 +97,15 @@ class EndpointHandler:
 
     def __call__(self, data: dict) -> dict:
         inputs = data.pop("inputs", data)
+        task = inputs.get("task", "generate_image").lower()
+
         model_variant = inputs.get("model_variant", "conditional").lower()
         pipeline_type = inputs.get("pipeline_type", model_variant).lower()
         conditions_input = inputs.get("conditions", ["No Finding"])
         guidance_scale = float(inputs.get("guidance_scale", 3.0))
         seed = int(inputs.get("seed", 42))
         num_inference_steps = int(inputs.get("num_inference_steps", 50))
-        p_mask = float(inputs.get("p_mask", 0.9)) 
+        p_mask = float(inputs.get("p_mask", 0.75)) 
 
         print(f"Request: model_variant='{model_variant}', pipeline_type='{pipeline_type}', conditions='{conditions_input}', guidance={guidance_scale}, seed={seed}, steps={num_inference_steps}")
 
@@ -145,36 +147,109 @@ class EndpointHandler:
             error_msg = f"Invalid pipeline_type: '{pipeline_type}'. Choose 'conditional' or 'ambient'."
             print(error_msg)
             return {"error": error_msg}
+        
+        if task == "tcnp_batch":
+            print("Processing task: tcnp_batch")
+            # Parameters for TCNP batch processing
+            start_index = int(inputs.get("start_index", 0))
+            num_to_process = int(inputs.get("num_conditions_to_process", 10)) # Chunk size
+            guidance_scale = float(inputs.get("guidance_scale", 1.5)) # Default from TCNP.py
+            num_inference_steps = int(inputs.get("num_inference_steps", 1000)) # Default from TCNP.py
+            base_seed = int(inputs.get("seed", 42)) # Default from TCNP.py
+
+            # --- Load labels.npy ---
+            try:
+                labels_file_path = os.path.join(self.model_dir, "labels.npy")
+                if not os.path.exists(labels_file_path):
+                    return {"error": "labels.npy not found in model directory.", "task_processed": task}
+                all_labels_np = np.load(labels_file_path)
+            except Exception as e:
+                return {"error": f"Failed to load labels.npy: {str(e)}", "task_processed": task}
+
+            # --- Determine the slice of labels to process for the current chunk ---
+            actual_end_index = min(start_index + num_to_process, all_labels_np.shape[0])
+            if start_index >= all_labels_np.shape[0]:
+                return {"error": f"start_index ({start_index}) is out of bounds for labels.npy (size: {all_labels_np.shape[0]}).", "task_processed": task}
+            
+            target_labels_np_slice = all_labels_np[start_index:actual_end_index]
+            actual_conditions_processed_count = target_labels_np_slice.shape[0]
+
+            if actual_conditions_processed_count == 0:
+                return {"error": "No conditions selected for processing with the given indices for this chunk.", "task_processed": task}
+
+            collected_tcnp_metrics_for_batch = []
+            for i in range(actual_conditions_processed_count):
+                current_condition_np = target_labels_np_slice[i]
+                # Convert numpy multi-hot vector to torch tensor
+                condition_vector = torch.from_numpy(current_condition_np).float().to(self.device)
+                class_labels_for_pipeline = condition_vector.unsqueeze(0) # Add batch dimension
+
+                current_iteration_seed = base_seed + (start_index + i) # Seed varies based on original index
+                generator = torch.Generator(device=self.device).manual_seed(current_iteration_seed)
+                
+                print(f"TCNP Batch: Processing condition (original index {start_index + i}) with seed {current_iteration_seed}...")
+                
+                with torch.no_grad():
+                    result = active_pipeline(
+                        batch_size=1, # Process one condition at a time from the batch
+                        generator=generator,
+                        num_inference_steps=num_inference_steps,
+                        output_type="pt", # We only need the TCNP tensor
+                        class_labels=class_labels_for_pipeline,
+                        guidance_scale=guidance_scale,
+                        mem=True # CRUCIAL: Enable TCNP calculation
+                    )
+                
+                if "TCNP" in result:
+                    # result["TCNP"] should be shape (1, num_inference_steps)
+                    collected_tcnp_metrics_for_batch.append(result["TCNP"].squeeze().cpu().numpy()) # Squeeze to (steps,)
+                else:
+                    print(f"Warning: TCNP not found in result for condition original index {start_index + i}")
+                    collected_tcnp_metrics_for_batch.append(np.zeros(num_inference_steps)) # Error placeholder
+
+            final_metrics_array_for_batch = np.array(collected_tcnp_metrics_for_batch)
+
+            return {
+                "task_processed": task,
+                "processed_batch_info": {
+                    "requested_start_index": start_index,
+                    "requested_num_to_process": num_to_process,
+                    "actual_conditions_processed_count": actual_conditions_processed_count,
+                    "original_indices_processed": list(range(start_index, start_index + actual_conditions_processed_count))
+                },
+                "aggregated_tcnp_metrics": final_metrics_array_for_batch.tolist()
+            }
 
-        condition_vector = torch.zeros(len(_CHEXPERT_CLASSES), device=self.device)
-        for cond_name in conditions_input:
-            if cond_name in _CHEXPERT_CLASSES:
-                condition_index = _CHEXPERT_CLASSES.index(cond_name)
-                condition_vector[condition_index] = 1.0
-        class_labels = condition_vector.unsqueeze(0).float()
-
-        generator = torch.Generator(device=self.device).manual_seed(seed)
-
-        print(f"Starting inference with {pipeline_type} pipeline using UNet '{model_variant}'...")
-        with torch.no_grad():
-            result = active_pipeline(
-                batch_size=1, generator=generator, num_inference_steps=num_inference_steps,
-                output_type="np", class_labels=class_labels, guidance_scale=guidance_scale
-            )
-        image_np = result["images"][0]
-        print("Inference complete.")
-
-        if image_np.ndim == 3 and image_np.shape[2] == 1: # Ensure it's grayscale C=1
-             image_np = np.repeat(image_np, 3, axis=2) # Convert to C=3 for PNG
-        elif image_np.ndim == 2: # Grayscale H,W
-            image_np = np.stack((image_np,)*3, axis=-1) # Convert to H,W,3
-
-        image_np = (image_np * 255).round().astype("uint8")
-        pil_image = Image.fromarray(image_np)
-
-        buffered = BytesIO()
-        pil_image.save(buffered, format="PNG")
-        image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        print("Image processed and encoded to base64.")
-
-        return {"generated_image_base64": image_base64, "model_variant_used": model_variant, "pipeline_used": pipeline_type}
+        else:
+            condition_vector = torch.zeros(len(_CHEXPERT_CLASSES), device=self.device)
+            for cond_name in conditions_input:
+                if cond_name in _CHEXPERT_CLASSES:
+                    condition_index = _CHEXPERT_CLASSES.index(cond_name)
+                    condition_vector[condition_index] = 1.0
+            class_labels = condition_vector.unsqueeze(0).float()
+
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+
+            print(f"Starting inference with {pipeline_type} pipeline using UNet '{model_variant}'...")
+            with torch.no_grad():
+                result = active_pipeline(
+                    batch_size=1, generator=generator, num_inference_steps=num_inference_steps,
+                    output_type="np", class_labels=class_labels, guidance_scale=guidance_scale
+                )
+            image_np = result["images"][0]
+            print("Inference complete.")
+
+            if image_np.ndim == 3 and image_np.shape[2] == 1: # Ensure it's grayscale C=1
+                image_np = np.repeat(image_np, 3, axis=2) # Convert to C=3 for PNG
+            elif image_np.ndim == 2: # Grayscale H,W
+                image_np = np.stack((image_np,)*3, axis=-1) # Convert to H,W,3
+
+            image_np = (image_np * 255).round().astype("uint8")
+            pil_image = Image.fromarray(image_np)
+
+            buffered = BytesIO()
+            pil_image.save(buffered, format="PNG")
+            image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            print("Image processed and encoded to base64.")
+
+            return {"generated_image_base64": image_base64, "model_variant_used": model_variant, "pipeline_used": pipeline_type}