handler.py

# Example: app.py for the Inference Endpoint
import os
import torch
from PIL import Image
import numpy as np
from diffusers import DDPMScheduler
import base64
from io import BytesIO

try:
    from cfg_diffusion import CustomClassConditionedUnet
    from conditional_ddpm_pipeline import ConditionalDDPMPipeline
    from ambient_diffusion import AmbientDDPMPipeline
    from constants import CHEXPERT_CLASSES as _CHEXPERT_CLASSES
except ImportError as e:
    print(f"Error importing from src: {e}. Ensure 'src' directory is structured correctly in the repo.")
    _CHEXPERT_CLASSES = [
        "No Finding", "Enlarged Cardiomediastinum", "Cardiomegaly", "Lung Opacity",
        "Lung Lesion", "Edema", "Consolidation", "Pneumonia", "Atelectasis",
        "Pneumothorax", "Pleural Effusion", "Pleural Other", "Fracture", "Support Devices"
    ]
    raise e


class EndpointHandler:
    def __init__(self, model_dir="."):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model_dir = model_dir
        self.loaded_unets_cache = {} # Cache for {variant_name: unet_model}
        self.loaded_schedulers_cache = {} # Cache for {variant_name: scheduler_instance}
        print(f"Initializing EndpointHandler on device: {self.device}")

        default_scheduler_path = os.path.join(self.model_dir, "scheduler")
        if os.path.exists(os.path.join(default_scheduler_path, "scheduler_config.json")):
            self.default_scheduler_config_path = default_scheduler_path
            print(f"Found default scheduler config at {default_scheduler_path}")
        else:
            self.default_scheduler_config_path = None
            print(f"Default scheduler config not found at {default_scheduler_path}. Will create default DDPMScheduler if needed.")
        
        print("EndpointHandler initialized.")

    def _check_unet_files_exist(self, unet_path: str) -> bool:
        config_path = os.path.join(unet_path, "config.json")
        weights_bin_path = os.path.join(unet_path, "diffusion_pytorch_model.bin")
        weights_safetensors_path = os.path.join(unet_path, "diffusion_pytorch_model.safetensors")
        return os.path.exists(config_path) and (os.path.exists(weights_bin_path) or os.path.exists(weights_safetensors_path))

    def _get_or_load_unet(self, model_variant_name: str):
        if model_variant_name in self.loaded_unets_cache:
            print(f"Using cached UNet for variant: {model_variant_name}")
            return self.loaded_unets_cache[model_variant_name]

        unet_load_path_primary = os.path.join(self.model_dir, f"unet_{model_variant_name}")
        unet_load_path_final = None

        if self._check_unet_files_exist(unet_load_path_primary):
            unet_load_path_final = unet_load_path_primary
        elif model_variant_name == "conditional": # Fallback for "conditional" to try "unet_cfg"
            unet_load_path_fallback = os.path.join(self.model_dir, "unet_cfg")
            if self._check_unet_files_exist(unet_load_path_fallback):
                unet_load_path_final = unet_load_path_fallback
            else:
                raise FileNotFoundError(f"UNet for variant '{model_variant_name}' not found at {unet_load_path_primary} or {unet_load_path_fallback}. Required: config.json and (diffusion_pytorch_model.bin or .safetensors)")
        else:
            raise FileNotFoundError(f"UNet for variant '{model_variant_name}' not found at {unet_load_path_primary}. Required: config.json and (diffusion_pytorch_model.bin or .safetensors)")

        print(f"Loading UNet for variant '{model_variant_name}' from {unet_load_path_final}...")
        unet = CustomClassConditionedUnet.from_pretrained(unet_load_path_final).to(self.device)
        unet.eval()
        self.loaded_unets_cache[model_variant_name] = unet
        print(f"Loaded and cached UNet for {model_variant_name} with in_channels: {unet.config.in_channels}")
        return unet

    def _get_or_load_scheduler(self, model_variant_name: str, pipeline_type: str):
        scheduler_key = f"scheduler_for_{pipeline_type}" 

        if scheduler_key in self.loaded_schedulers_cache:
            print(f"Using cached scheduler configured for {pipeline_type}")
            return self.loaded_schedulers_cache[scheduler_key]

        if self.default_scheduler_config_path:
            scheduler = DDPMScheduler.from_pretrained(self.default_scheduler_config_path)
            print(f"Loaded base scheduler from {self.default_scheduler_config_path}")
        else:
            scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")
            print("Created new default DDPMScheduler instance.")

        if model_variant_name == "ambient":
            scheduler.config.prediction_type = "sample"
        else: 
            scheduler.config.prediction_type = "epsilon"
        
        print(f"Scheduler configured with prediction_type: '{scheduler.config.prediction_type}' for {pipeline_type} pipeline.")
        self.loaded_schedulers_cache[scheduler_key] = scheduler
        return scheduler

    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.75)) 

        print(f"Request: model_variant='{model_variant}', pipeline_type='{pipeline_type}', conditions='{conditions_input}', guidance={guidance_scale}, seed={seed}, steps={num_inference_steps}")

        try:
            unet_to_use = self._get_or_load_unet(model_variant)
            scheduler_to_use = self._get_or_load_scheduler(model_variant, pipeline_type)
        except FileNotFoundError as e:
            print(f"Error loading model components: {e}")
            return {"error": str(e)}
        except Exception as e:
            print(f"Unexpected error during model/scheduler loading: {e}")
            return {"error": f"Unexpected error loading components: {str(e)}"}

        unet_in_channels = unet_to_use.config.in_channels
        active_pipeline = None

        if pipeline_type == "conditional":
            if unet_in_channels != 1:
                error_msg = f"UNet incompatible: Requested conditional pipeline, but loaded UNet '{model_variant}' has {unet_in_channels} input channels (expected 1)."
                print(error_msg)
                return {"error": error_msg}
            if scheduler_to_use.config.prediction_type != "epsilon":
                 print(f"Warning: Conditional pipeline expects 'epsilon' prediction, but scheduler is '{scheduler_to_use.config.prediction_type}'. Forcing to 'epsilon'.")
                 scheduler_to_use.config.prediction_type = "epsilon" 
            active_pipeline = ConditionalDDPMPipeline(unet=unet_to_use, scheduler=scheduler_to_use)
            print(f"Using ConditionalDDPMPipeline with UNet '{model_variant}'.")

        elif pipeline_type == "ambient":
            if unet_in_channels != 2:
                error_msg = f"UNet incompatible: Requested ambient pipeline, but loaded UNet '{model_variant}' has {unet_in_channels} input channels (expected 2 for image+mask)."
                print(error_msg)
                return {"error": error_msg}
            if scheduler_to_use.config.prediction_type != "sample":
                print(f"Warning: Ambient pipeline expects 'sample' prediction, but scheduler is '{scheduler_to_use.config.prediction_type}'. Forcing to 'sample'.")
                scheduler_to_use.config.prediction_type = "sample" 
            active_pipeline = AmbientDDPMPipeline(unet=unet_to_use, scheduler=scheduler_to_use, p_mask=p_mask)
            print(f"Using AmbientDDPMPipeline with UNet '{model_variant}' and p_mask: {p_mask}.")
        else:
            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()
            }

        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}