new layout

app.py

@@ -11,24 +11,16 @@ import torch
 from PIL import Image
 from torchvision import transforms
 
-from evals import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDataset, Metric
+from factories import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDataset, Metric
 
 
-DEVICE_STR = 'cuda'
+### Config
+DEVICE_STR = 'cuda'            # run model inference on NVIDIA gpu
+torch.set_grad_enabled(False)  # stops tracking values for gradients
 
 
 ### Gradio Utils
 
-def generate_imgs_from_dataset(dataset: EvalDataset, idx: int,
-                               model: EvalModel, baseline: BaselineModel,
-                               physics: PhysicsWithGenerator, use_gen: bool,
-                               metrics: List[Metric]):
-    ### Load 1 image
-    x = dataset[idx]    # shape : (3, 256, 256)
-    x = x.unsqueeze(0)  # shape : (1, 3, 256, 256)
-
-    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
-
 def generate_imgs_from_user(image,
                             model: EvalModel, baseline: BaselineModel,
                             physics: PhysicsWithGenerator, use_gen: bool,
@@ -37,9 +29,31 @@ def generate_imgs_from_user(image,
         return None, None, None, None, None, None, None, None
 
     # PIL image -> torch.Tensor
-    x = transforms.ToTensor()(image).unsqueeze(0).to('cuda')
+    x = transforms.ToTensor()(image).unsqueeze(0).to(DEVICE_STR)
+
+    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+
+def generate_imgs_from_dataset(dataset: EvalDataset, idx: int,
+                               model: EvalModel, baseline: BaselineModel,
+                               physics: PhysicsWithGenerator, use_gen: bool,
+                               metrics: List[Metric]):
+    ### Load 1 image
+    x = dataset[idx]    # shape : (C, H, W)
+    x = x.unsqueeze(0)  # shape : (1, C, H, W)
 
     return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+
+def generate_random_imgs_from_dataset(dataset: EvalDataset,
+                                      model: EvalModel,
+                                      baseline: BaselineModel,
+                                      physics: PhysicsWithGenerator,
+                                      use_gen: bool,
+                                      metrics: List[Metric]):
+    idx = random.randint(0, len(dataset)-1)
+    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs_from_dataset(
+        dataset, idx, model, baseline, physics, use_gen, metrics
+        )
+    return idx, x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline
     
 def generate_imgs(x: torch.Tensor,
                   model: EvalModel, baseline: BaselineModel,
@@ -75,7 +89,7 @@ def generate_imgs(x: torch.Tensor,
             metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
 
     ### Process y when y shape is different from x shape
-    if physics.name == "MRI" or "SR" in physics.name:
+    if physics.name == "MRI" in physics.name:
         y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
     else:
         y_plot = y.clone()
@@ -93,18 +107,6 @@ def generate_imgs(x: torch.Tensor,
 
     return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
-def generate_random_imgs_from_dataset(dataset: EvalDataset,
-                                        model: EvalModel,
-                                        baseline: BaselineModel,
-                                        physics: PhysicsWithGenerator,
-                                        use_gen: bool,
-                                        metrics: List[Metric]):
-    idx = random.randint(0, len(dataset)-1)
-    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs_from_dataset(
-        dataset, idx, model, baseline, physics, use_gen, metrics
-        )
-    return idx, x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline
-
 
 get_list_metrics_on_DEVICE_STR = partial(Metric.get_list_metrics, device_str=DEVICE_STR)
 get_eval_model_on_DEVICE_STR = partial(EvalModel, device_str=DEVICE_STR)
@@ -112,7 +114,8 @@ get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
 get_physics_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
 
-AVAILABLE_PHYSICS = PhysicsWithGenerator.all_physics
+AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
+                     'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
 def get_dataset(dataset_name):
     global AVAILABLE_PHYSICS
     if dataset_name == 'MRI':
@@ -124,10 +127,15 @@ def get_dataset(dataset_name):
         baseline_name = 'DPIR_CT'
         physics_name = 'CT'
     else:
-        AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard', 'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
+        AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
+                             'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
         baseline_name = 'DPIR'
         physics_name = 'MotionBlur_easy'
-    return get_dataset_on_DEVICE_STR(dataset_name), get_physics_on_DEVICE_STR(physics_name), get_baseline_model_on_DEVICE_STR(baseline_name)
+
+    dataset = get_dataset_on_DEVICE_STR(dataset_name)
+    physics = get_physics_on_DEVICE_STR(physics_name)
+    baseline = get_baseline_model_on_DEVICE_STR(baseline_name)
+    return dataset, physics, baseline
 
 
 ### Gradio Blocks interface
@@ -135,9 +143,9 @@ def get_dataset(dataset_name):
 # Define custom CSS
 custom_css = """
 .fixed-textbox textarea {
-    height: 90px !important;  /* Adjust height to fit exactly 4 lines */
-    overflow: scroll;         /* Add a scroll bar if necessary */
-    resize: none;             /* User can resize vertically the textbox */
+    height: 100px !important;  /* Adjust height to fit exactly 4 lines */
+    overflow: scroll;          /* Add a scroll bar if necessary */
+    resize: none;              /* User can resize vertically the textbox */
 }
 """
 
@@ -145,87 +153,88 @@ title = "Inverse problem playground"  # displayed on gradio tab and in the gradi
 with gr.Blocks(title=title, css=custom_css) as interface:
     gr.Markdown("## " + title)
 
-    # Loading things
-    model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))  # lambda expression to instanciate a callable in a gr.State
-    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))  # lambda expression to instanciate a callable in a gr.State
-    dataset_placeholder = gr.State(lambda: get_dataset_on_DEVICE_STR("Natural"))
-    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))  # lambda expression to instanciate a callable in a gr.State
-    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))
+    # DEFAULT VALUES
+    # Issue: giving directly a `torch.nn.module` to `gr.State(...)` since it has __call__ method
+    # Solution: using lambda expression
+    model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))
+    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
+    dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
+    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
+    idx_placeholder = gr.State(0)
+
+    metric_names = ["PSNR"]
+    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(metric_names))
 
-    @gr.render(inputs=[dataset_placeholder, physics_placeholder, metrics_placeholder])
-    def dynamic_layout(dataset, physics, metrics):
+    @gr.render(inputs=[dataset_placeholder, physics_placeholder])
+    def dynamic_layout(dataset, physics):
         ### LAYOUT
-        dataset_name = dataset.name
-        physics_name = physics.name
-        metric_names = [metric.name for metric in metrics]
 
-        # Components: Inputs/Outputs + Load EvalDataset/PhysicsWithGenerator/EvalModel/BaselineModel
+        # Display images
+        with gr.Row():
+            gt_img = gr.Image(label=f"Ground-truth IMAGE", interactive=True)
+            observed_img = gr.Image(label=f"Observed IMAGE", interactive=False)
+            model_a_out = gr.Image(label="RAM OUTPUT", interactive=False)
+            model_b_out = gr.Image(label="DPIR OUTPUT", interactive=False)
+
+        # Manage datasets and display metric values
         with gr.Row():
             with gr.Column():
+                run_button = gr.Button("Demo on above image")
+                choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
+                                          label="Datasets",
+                                          value=dataset.name)
+                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index")
                 with gr.Row():
-                    with gr.Column():
-                        clean = gr.Image(label=f"{dataset_name} IMAGE", interactive=True)
-                        physics_params = gr.Textbox(label="Physics parameters", elem_classes=["fixed-textbox"], value=physics.display_saved_params())
-                    with gr.Column():
-                        y_image = gr.Image(label=f"{physics_name} IMAGE", interactive=False)
-                        y_metrics = gr.Textbox(label="Metrics(y, x)", elem_classes=["fixed-textbox"],)
+                    load_button = gr.Button("Run on index image from dataset")
+                    load_random_button = gr.Button("Run on random image from dataset")
+            with gr.Column():
+                observed_metrics = gr.Textbox(label="PSNR(Observed, Ground-truth)",
+                                              elem_classes=["fixed-textbox"])
+            with gr.Column():
+                out_a_metric = gr.Textbox(label="PSNR(RAM(Observed, Ground-truth)",
+                                          elem_classes=["fixed-textbox"])
+            with gr.Column():
+                out_b_metric = gr.Textbox(label="PSNR(DPIR(Observed, Ground-truth)",
+                                          elem_classes=["fixed-textbox"])
 
+        # Manage physics
+        with gr.Row():
+            with gr.Column(scale=1):
                 choose_physics = gr.Radio(choices=AVAILABLE_PHYSICS,
-                                          label="List of PhysicsWithGenerator",
-                                          value=physics_name)
+                                          label="Physics",
+                                          value=physics.name)
+                use_generator_button = gr.Checkbox(label="Generate physics parameters during inference")
+            with gr.Column(scale=1):
                 with gr.Row():
                     key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()),
-                                               label="Updatable Parameter Key",
-                                               scale=2)
-                    value_text = gr.Textbox(label="Update Value", scale=2)
-                    update_button = gr.Button("Manually update parameter value", scale=1)
-
-            with gr.Column():
-                with gr.Row():
-                    with gr.Column():
-                        model_a_out = gr.Image(label="RAM OUTPUT", interactive=False)
-                        out_a_metric = gr.Textbox(label="Metrics(RAM(y, physics), x)", elem_classes=["fixed-textbox"])
-                    with gr.Column():
-                        model_b_out = gr.Image(label="DPIR OUTPUT", interactive=False)
-                        out_b_metric = gr.Textbox(label="Metrics(DPIR(y, physics), x)", elem_classes=["fixed-textbox"])
-                with gr.Row():
-                    choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
-                                              label="List of EvalDataset",
-                                              value=dataset_name,
-                                              scale=2)
-                    idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", scale=1)
+                                                label="Updatable Parameter Key")
+                    value_text = gr.Textbox(label="Update Value")
+                update_button = gr.Button("Manually update parameter value")
+            with gr.Column(scale=2):
+                physics_params = gr.Textbox(label="Physics parameters",
+                                            elem_classes=["fixed-textbox"],
+                                            value=physics.display_saved_params())  
 
-        # Components: Load Metric + Load image Buttons
-        with gr.Row():
-            with gr.Column(scale=3):
-                choose_metrics = gr.CheckboxGroup(choices=Metric.all_metrics,
-                                                  value=metric_names,
-                                                  label="Choose metrics you are interested")
-            use_generator_button = gr.Checkbox(label="Generate valid physics parameters", scale=1)
-            run_button = gr.Button("Run current image", scale=1)
-            with gr.Column(scale=1):
-                load_button = gr.Button("Load images from dataset...")
-                load_random_button = gr.Button("Load randomly from dataset...")
 
         ### Event listeners
+
         choose_dataset.change(fn=get_dataset,
                               inputs=choose_dataset,
                               outputs=[dataset_placeholder, physics_placeholder, model_b_placeholder])
         choose_physics.change(fn=get_physics_on_DEVICE_STR,
                               inputs=choose_physics,
                               outputs=[physics_placeholder])
-        update_button.click(fn=physics.update_and_display_params, inputs=[key_selector, value_text], outputs=physics_params)
-        choose_metrics.change(fn=get_list_metrics_on_DEVICE_STR,
-                              inputs=choose_metrics,
-                              outputs=metrics_placeholder)
+        update_button.click(fn=physics.update_and_display_params,
+                            inputs=[key_selector, value_text], outputs=physics_params)
         run_button.click(fn=generate_imgs_from_user,
-                         inputs=[clean,
+                         inputs=[gt_img,
                                  model_a_placeholder,
                                  model_b_placeholder,
                                  physics_placeholder,
                                  use_generator_button,
                                  metrics_placeholder],
-                         outputs=[clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
+                         outputs=[gt_img, observed_img, model_a_out, model_b_out,
+                                  physics_params, observed_metrics, out_a_metric, out_b_metric])
         load_button.click(fn=generate_imgs_from_dataset,
                           inputs=[dataset_placeholder,
                                   idx_slider,
@@ -234,7 +243,8 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                                   physics_placeholder,
                                   use_generator_button,
                                   metrics_placeholder],
-                          outputs=[clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
+                          outputs=[gt_img, observed_img, model_a_out, model_b_out,
+                                   physics_params, observed_metrics, out_a_metric, out_b_metric])
         load_random_button.click(fn=generate_random_imgs_from_dataset,
                                  inputs=[dataset_placeholder,
                                          model_a_placeholder,
@@ -242,6 +252,8 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                                          physics_placeholder,
                                          use_generator_button,
                                          metrics_placeholder],
-                                 outputs=[idx_slider, clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
+                                 outputs=[idx_slider, gt_img, observed_img, model_a_out, model_b_out,
+                                          physics_params, observed_metrics, out_a_metric, out_b_metric])
+
 
 interface.launch()

evals.py → factories.py

@@ -8,6 +8,7 @@ from torchvision import transforms
 
 from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI, LsdirMiniDataset
 from model_factory import get_model
+from physics.blur_generator import GaussianBlurGenerator
 
 DEFAULT_MODEL_PARAMS = {
     "in_channels": [1, 2, 3],
@@ -159,7 +160,7 @@ class PhysicsWithGenerator(torch.nn.Module):
 
     def _update_save_params(self, key: str, value: Any) -> None:
         """Update value of an existing key in save_params."""
-        if key in list(self.saved_params["updatable_params"].keys()):
+        if value != "" and key in list(self.saved_params["updatable_params"].keys()):
             if type(value) == str:  # it may be only a str representation
                 # type: str -> ???
                 value = self.saved_params["updatable_params_converter"][key](value)
@@ -168,7 +169,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                 value = float(f"{value:.4f}")  # keeps only 4 significant digits
             self.saved_params["updatable_params"][key] = value
 
-    def update_and_display_params(self, key, value) -> str:
+    def update_and_display_params(self, key: str, value: Any) -> str:
         """_update_save_params + update physics with saved_params + display_saved_params"""
         self._update_save_params(key, value)
 

models/unrolled_dpir.py

@@ -1,304 +0,0 @@
-import numpy as np
-import deepinv
-import torch
-import deepinv as dinv
-from deepinv.optim.data_fidelity import L2
-from deepinv.optim.prior import PnP
-from deepinv.unfolded import unfolded_builder
-import copy
-import deepinv.optim.utils
-
-class PoissonGaussianDistance(dinv.optim.Distance):
-    r"""
-    Implementation of :math:`\distancename` as the normalized :math:`\ell_2` norm
-
-    .. math::
-        f(x) = (x-y)^{T}\Sigma_y(x-y)
-
-    with :math:`\Sigma_y=\text{diag}(gamma y + \sigma^2)`
-
-    :param float sigma: Gaussian noise parameter. Default: 1.
-    :param float gain: Poisson noise parameter. Default 0.
-    """
-
-    def __init__(self, sigma=1.0, gain=0.):
-        super().__init__()
-        self.sigma = sigma
-        self.gain = gain
-
-    def fn(self, x, y, *args, **kwargs):
-        r"""
-        Computes the distance :math:`\distance{x}{y}` i.e.
-
-        .. math::
-
-            \distance{x}{y} = \frac{1}{2}\|x-y\|^2
-
-
-        :param torch.Tensor u: Variable :math:`x` at which the data fidelity is computed.
-        :param torch.Tensor y: Data :math:`y`.
-        :return: (:class:`torch.Tensor`) data fidelity :math:`\datafid{u}{y}` of size `B` with `B` the size of the batch.
-        """
-        norm = 1.0 / (self.sigma**2 + y * self.gain)
-        z = (x - y) * norm
-        d = 0.5 * torch.norm(z.reshape(z.shape[0], -1), p=2, dim=-1) ** 2
-        return d
-
-    def grad(self, x, y, *args, **kwargs):
-        r"""
-        Computes the gradient of :math:`\distancename`, that is  :math:`\nabla_{x}\distance{x}{y}`, i.e.
-
-        .. math::
-
-            \nabla_{x}\distance{x}{y} = \frac{1}{\sigma^2} x-y
-
-
-        :param torch.Tensor x: Variable :math:`x` at which the gradient is computed.
-        :param torch.Tensor y: Observation :math:`y`.
-        :return: (:class:`torch.Tensor`) gradient of the distance function :math:`\nabla_{x}\distance{x}{y}`.
-        """
-        norm = 1.0 / (self.sigma**2 + y * self.gain)
-        return (x - y) * norm
-
-    def prox(self, x, y, *args, gamma=1.0, **kwargs):
-        r"""
-        Proximal operator of :math:`\gamma \distance{x}{y} = \frac{\gamma}{2 \sigma^2} \|x-y\|^2`.
-
-        Computes :math:`\operatorname{prox}_{\gamma \distancename}`, i.e.
-
-        .. math::
-
-           \operatorname{prox}_{\gamma \distancename} = \underset{u}{\text{argmin}} \frac{\gamma}{2\sigma^2}\|u-y\|_2^2+\frac{1}{2}\|u-x\|_2^2
-
-
-        :param torch.Tensor x: Variable :math:`x` at which the proximity operator is computed.
-        :param torch.Tensor y: Data :math:`y`.
-        :param float gamma: thresholding parameter.
-        :return: (:class:`torch.Tensor`) proximity operator :math:`\operatorname{prox}_{\gamma \distancename}(x)`.
-        """
-        norm = 1.0 / (self.sigma**2 + y * self.gain)
-        return (x + norm * gamma * y) / (1 + gamma * norm)
-
-
-class PoissonGaussianDataFidelity(dinv.optim.DataFidelity):
-    r"""
-    Implementation of the data-fidelity as the normalized :math:`\ell_2` norm
-
-    .. math::
-
-        f(x) = \|\forw{x}-y\|^2_{\text{diag}(\sigma^2 + y \gamma)}
-
-    It can be used to define a log-likelihood function associated with Poisson Gaussian noise
-    by setting an appropriate noise level :math:`\sigma`.
-
-    :param float sigma: Standard deviation of the noise to be used as a normalisation factor.
-    :param float gain: Gain factor of the data-fidelity term.
-    """
-
-    def __init__(self, sigma=1.0, gain=0.):
-        super().__init__()
-        self.d = PoissonGaussianDistance(sigma=sigma, gain=gain)
-        self.gain = gain
-        self.sigma = sigma
-
-    def prox(self, x, y, physics, gamma=1.0, *args, **kwargs):
-        r"""
-        Proximal operator of :math:`\gamma \datafid{Ax}{y} = \frac{\gamma}{2\sigma^2}\|Ax-y\|^2`.
-
-        Computes :math:`\operatorname{prox}_{\gamma \datafidname}`, i.e.
-
-        .. math::
-
-           \operatorname{prox}_{\gamma \datafidname} = \underset{u}{\text{argmin}} \frac{\gamma}{2\sigma^2}\|Au-y\|_2^2+\frac{1}{2}\|u-x\|_2^2
-
-
-        :param torch.Tensor x: Variable :math:`x` at which the proximity operator is computed.
-        :param torch.Tensor y: Data :math:`y`.
-        :param deepinv.physics.Physics physics: physics model.
-        :param float gamma: stepsize of the proximity operator.
-        :return: (:class:`torch.Tensor`) proximity operator :math:`\operatorname{prox}_{\gamma \datafidname}(x)`.
-        """
-        assert isinstance(physics, dinv.physics.LinearPhysics), "not implemented for non-linear physics"   
-        if isinstance(physics, dinv.physics.StackedPhysics):
-            device=y[0].device
-            noise_model = physics[-1].noise_model
-        else:
-            device=y.device
-            noise_model = physics.noise_model
-        if hasattr(noise_model, "gain"):
-            self.gain = noise_model.gain.detach().to(device)
-        if hasattr(noise_model, "sigma"):
-            self.sigma = noise_model.sigma.detach().to(device)
-        # Ensure sigma is a tensor and reshape if necessary
-        if isinstance(self.sigma, float):
-            self.sigma = torch.tensor([self.sigma], device=device)
-        if self.sigma.ndim == 0 :
-            self.sigma = self.sigma.unsqueeze(0).to(device)
-        # Ensure gain is a tensor and reshape if necessary
-        if isinstance(self.gain, float):
-            self.gain = torch.tensor([self.gain], device=device)
-        if self.gain.ndim == 0 :
-            self.gain = self.gain.unsqueeze(0).to(device)
-        if self.gain[0] > 0 :
-            norm = gamma / (self.sigma[:, None, None, None]**2 + y * self.gain[:, None, None, None])
-        else : 
-            norm = gamma / (self.sigma[:, None, None, None]**2)
-        A = lambda u: physics.A_adjoint(physics.A(u)*norm) + u
-        b = physics.A_adjoint(norm*y) + x
-        return deepinv.optim.utils.conjugate_gradient(A, b, init=x, max_iter=3, tol=1e-3)
-
-from deepinv.optim.optim_iterators import OptimIterator, fStep, gStep
-
-class myHQSIteration(OptimIterator):
-    r"""
-    Single iteration of half-quadratic splitting.
-
-    Class for a single iteration of the Half-Quadratic Splitting (HQS) algorithm for minimising :math:`f(x) + \lambda \regname(x)`.
-    The iteration is given by
-
-
-    .. math::
-        \begin{equation*}
-        \begin{aligned}
-        u_{k} &= \operatorname{prox}_{\gamma f}(x_k) \\
-        x_{k+1} &= \operatorname{prox}_{\sigma \lambda \regname}(u_k).
-        \end{aligned}
-        \end{equation*}
-
-
-    where :math:`\gamma` and :math:`\sigma` are step-sizes. Note that this algorithm does not converge to
-    a minimizer of :math:`f(x) + \lambda  \regname(x)`, but instead to a minimizer of
-    :math:`\gamma\, ^1f+\sigma \lambda \regname`, where :math:`^1f` denotes
-    the Moreau envelope of :math:`f`
-
-    """
-
-    def __init__(self, **kwargs):
-        super(myHQSIteration, self).__init__(**kwargs)
-        self.g_step = mygStepHQS(**kwargs)
-        self.f_step = myfStepHQS(**kwargs)
-        self.requires_prox_g = True
-
-class myfStepHQS(fStep):
-    r"""
-    HQS fStep module.
-    """
-
-    def __init__(self, **kwargs):
-        super(myfStepHQS, self).__init__(**kwargs)
-
-    def forward(self, x, cur_data_fidelity, cur_params, y, physics):
-        r"""
-        Single proximal step on the data-fidelity term :math:`f`.
-
-        :param torch.Tensor x: Current iterate :math:`x_k`.
-        :param deepinv.optim.DataFidelity cur_data_fidelity: Instance of the DataFidelity class defining the current data_fidelity.
-        :param dict cur_params: Dictionary containing the current parameters of the algorithm.
-        :param torch.Tensor y: Input data.
-        :param deepinv.physics.Physics physics: Instance of the physics modeling the data-fidelity term.
-        """
-        return cur_data_fidelity.prox(x, y, physics, gamma=cur_params["stepsize"])
-
-class mygStepHQS(gStep):
-    r"""
-    HQS gStep module.
-    """
-
-    def __init__(self, **kwargs):
-        super(mygStepHQS, self).__init__(**kwargs)
-
-    def forward(self, x, cur_prior, cur_params):
-        r"""
-        Single proximal step on the prior term :math:`\lambda \regname`.
-
-        :param torch.Tensor x: Current iterate :math:`x_k`.
-        :param dict cur_prior: Class containing the current prior.
-        :param dict cur_params: Dictionary containing the current parameters of the algorithm.
-        """
-        return cur_prior.prox(
-            x,
-            sigma_denoiser = cur_params["g_param"],
-            gain_denoiser = cur_params["gain_param"],
-            gamma=cur_params["lambda"] * cur_params["stepsize"],
-        )
-
-
-def get_unrolled_architecture(gain_param_init = 1e-3, weight_tied = True, model = None, device = 'cpu'):
-
-    # Unrolled optimization algorithm parameters
-    max_iter = 8  # number of unfolded layers
-
-    # Select the data fidelity term
-    
-
-    # Set up the trainable denoising prior
-    # Here the prior model is common for all iterations
-    if model is not None : 
-        denoiser = model.to(device)
-    else :
-        denoiser = dinv.models.DRUNet(
-        pretrained= '/lustre/fswork/projects/rech/nyd/commun/mterris/base_checkpoints/drunet_deepinv_color_finetune_22k.pth',
-        ).to(device)
-
-    class myPnP(PnP):
-        r"""
-        Gradient-Step Denoiser prior.
-        """
-
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-
-        def prox(self, x, sigma_denoiser, gain_denoiser, *args, **kwargs):
-            if not self.training:
-                pad = (-x.size(-2) % 8, -x.size(-1) % 8)
-                x = torch.nn.functional.pad(x, (0, pad[1], 0, pad[0]), mode="constant")
-            out = self.denoiser(x, sigma=sigma_denoiser, gamma=gain_denoiser)
-            if not self.training:
-                out = out[..., : -pad[0] or None, : -pad[1] or None]
-            return out
-    
-    data_fidelity = PoissonGaussianDataFidelity()
-    
-    if not weight_tied :
-        prior = [myPnP(denoiser=copy.deepcopy(denoiser)) for i in range(max_iter)]
-    else :
-        prior = [myPnP(denoiser=denoiser)]
-    
-    def get_DPIR_params(noise_level_img, max_iter=8):
-        r"""
-        Default parameters for the DPIR Plug-and-Play algorithm.
-
-        :param float noise_level_img: Noise level of the input image.
-        """
-        s1 = 49.0 / 255.0
-        s2 = noise_level_img
-        sigma_denoiser = np.logspace(np.log10(s1), np.log10(s2), max_iter).astype(
-            np.float32
-        )
-        stepsize = (sigma_denoiser / max(0.01, noise_level_img)) ** 2
-        lamb = 1 / 0.23
-        return list(sigma_denoiser), list(lamb * stepsize)
-
-    sigma_denoiser, stepsize = get_DPIR_params(0.05)
-    stepsize = torch.tensor(stepsize) * (torch.tensor(sigma_denoiser)**2)
-    gain_denoiser = [gain_param_init]*len(sigma_denoiser)
-    params_algo = {"stepsize": stepsize, "g_param": sigma_denoiser, "gain_param": gain_denoiser}
-
-    trainable_params = [
-        "g_param",
-        "gain_param"
-        "stepsize",
-    ]  # define which parameters from 'params_algo' are trainable
-
-    # Define the unfolded trainable model.
-    model = unfolded_builder(
-        iteration=myHQSIteration(),
-        params_algo=params_algo.copy(),
-        trainable_params=trainable_params,
-        data_fidelity=data_fidelity,
-        max_iter=max_iter,
-        prior=prior,
-        device=device,
-    )
-
-    return model.to(device)