app.py

# app.py
import gradio as gr
import os

# Import the core logic from your "normal inference file"
from NoiseFilter.MaIR.inference_runner import MaIR_Upsampler

# --- Global model cache for performance ---
# This dictionary will store loaded models to avoid reloading on every API call.
model_cache = {}

def get_model(model_name):
    """Loads a model into the cache if it's not already there."""
    if model_name not in model_cache:
        print(f"Loading model {model_name} into cache...")
        model_cache[model_name] = MaIR_Upsampler(model_name=model_name)
    return model_cache[model_name]

# --- API Function ---
def inference_api(image, model_name):
    """
    This is the function that the API will call.
    It takes a NumPy array and a model name string as input.
    """
    if image is None:
        # Gradio handles this by not running, but good practice for raw API calls.
        raise ValueError("No image provided.")
    
    upsampler = get_model(model_name)
    output_image = upsampler.process(image)
    return output_image

# --- Create the Gradio Interface (for API generation) ---
# We define a minimal interface. The primary goal is API exposure.
interface = gr.Interface(
    fn=inference_api,
    inputs=[
        gr.Image(type="numpy", label="Input Image"),
        gr.Dropdown(
            choices=['MaIR-SRx4', 'MaIR-SRx2', 'MaIR-CDN-s50'], 
            value='MaIR-SRx4', 
            label="Select Model"
        ),
    ],
    outputs=gr.Image(type="numpy", label="Output Image"),
    title="MaIR: Image Restoration API",
    description="API for MaIR models. Use the '/api' endpoint for programmatic access."
)

# Launch the app. This will start the web server and create the API.
interface.launch()