# 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}