<!-- After Hugging Face Login install these libraries -->
<!-- !pip install torchsr
!pip install diffusers[training] -->
from torchsr.models import edsr
from diffusers import DDPMPipeline
import torch
from PIL import Image
import torchvision.transforms as transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load the butterfly pipeline
butterfly_pipeline = DDPMPipeline.from_pretrained(
    "Ketansomewhere/Lung_Ultrasound_Diffusion_720p"
).to(device)

# Create 1(can be n in principle) images
images = butterfly_pipeline(batch_size=1).images

# Load the pre-trained EDSR model
model = edsr(scale=4, pretrained=True).to(device)

upscaled_images = []

for img in images:
    # Convert to tensor and add batch dimension
    img_tensor = transforms.ToTensor()(img).unsqueeze(0).to(device)

    # Upscale the image
    upscaled_img_tensor = model(img_tensor)

    # Remove batch dimension and convert back to PIL image
    upscaled_img = transforms.ToPILImage()(upscaled_img_tensor.squeeze(0).cpu())

    # Add to list of upscaled images
    upscaled_images.append(upscaled_img)

# Function to make a grid of images
def make_grid(images, size=720):
    """Given a list of PIL images, stack them together into a line for easy viewing"""
    output_im = Image.new("RGB", (size * len(images), size))
    for i, im in enumerate(images):
        output_im.paste(im.resize((size, size)), (i * size, 0))
    return output_im

# View the result
make_grid(upscaled_images)