Merge branch 'main' of https://huggingface.co/isp-uv-es/superXI

ldm-baseline/metadata.json

@@ -1,10 +0,0 @@
-{
-    "name": "ldm-baseline",
-    "authors": ["CompVis team"],
-    "affiliations": ["None"],
-    "description": "A baseline of LDM models trained on the Open Images dataset.",
-    "code": "open-source",
-    "scale": "x4",
-    "url": "https://huggingface.co/CompVis/ldm-super-resolution-4x-openimages",
-    "license": "apache-2.0"
-  }

ldm-baseline/run.py

@@ -1,30 +0,0 @@
-from diffuser.utils import create_stable_diffusion_model, run_diffuser
-import opensr_test
-import matplotlib.pyplot as plt
-
-# Load the model
-model = create_stable_diffusion_model(device="cuda")
-
-# Load the dataset
-dataset = opensr_test.load("naip")
-lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]
-
-# Run the model
-results = run_diffuser(
-    model=model,
-    lr=lr_dataset[5][:,0:64, 0:64],
-    hr=hr_dataset[5][:,0:256, 0:256],
-    device="cuda"
-)
-
-# Display the results
-fig, ax = plt.subplots(1, 3, figsize=(10, 5))
-ax[0].imshow(results["lr"].transpose(1, 2, 0)/3000)
-ax[0].set_title("LR")
-ax[0].axis("off")
-ax[1].imshow(results["sr"].transpose(1, 2, 0)/3000)
-ax[1].set_title("SR")
-ax[1].axis("off")
-ax[2].imshow(results["hr"].transpose(1, 2, 0) / 3000)
-ax[2].set_title("HR")
-plt.show()

ldm-baseline/utils.py

@@ -1,81 +0,0 @@
-from diffusers import LDMSuperResolutionPipeline
-import numpy as np
-import opensr_test
-import torch
-import pickle
-from typing import Union
-
-
-def create_stable_diffusion_model(
-    device: Union[str, torch.device] = "cuda"
-) -> LDMSuperResolutionPipeline:
-    """ Create the stable diffusion model
-
-    Returns:
-        LDMSuperResolutionPipeline: The model to use for
-            super resolution.
-    """
-    model_id = "CompVis/ldm-super-resolution-4x-openimages"
-    pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)
-    pipeline = pipeline.to(device)
-    return pipeline
-
-def run_diffuser(
-    model: LDMSuperResolutionPipeline,
-    lr: torch.Tensor,
-    hr: torch.Tensor,
-    device: Union[str, torch.device] = "cuda"    
-) -> dict:
-    """ Run the model on the low resolution image
-
-    Args:
-        model (LDMSuperResolutionPipeline): The model to use
-        lr (torch.Tensor): The low resolution image
-        hr (torch.Tensor): The high resolution image
-        device (Union[str, torch.device], optional): The device
-            to use. Defaults to "cuda".
-
-    Returns:
-        dict: The results of the model
-    """
-
-    # move the images to the device
-    lr = (torch.from_numpy(lr[[3, 2, 1]]) / 2000).to(device).clamp(0, 1)
-    
-    if lr.shape[1] == 121:
-        # add padding
-        lr = torch.nn.functional.pad(
-            lr[None],
-            pad=(3, 4, 3, 4),
-            mode='reflect'
-        ).squeeze()
-    
-        # run the model
-        with torch.no_grad():
-            sr = model(lr[None], num_inference_steps=100, eta=1)
-            sr = torch.from_numpy(
-                np.array(sr.images[0])/255
-            ).permute(2,0,1).float()
-
-        # remove padding
-        sr = sr[:, 3*4:-4*4, 3*4:-4*4]
-        lr = lr[:, 3:-4, 3:-4]
-    else:
-        # run the model
-        with torch.no_grad():
-            sr = model(lr[None], num_inference_steps=100, eta=1)
-            sr = torch.from_numpy(
-                np.array(sr.images[0])/255
-            ).permute(2,0,1).float()
-
-    lr = (lr.cpu().numpy() * 2000).astype(np.uint16)
-    hr = ((hr[0:3] / 2000).clip(0, 1) * 2000).astype(np.uint16)
-    sr = (sr.cpu().numpy() * 2000).astype(np.uint16)
-
-    results = {
-        "lr": lr,
-        "hr": hr,
-        "sr": sr
-    }
-    
-    return results