from diffusers import LDMSuperResolutionPipeline
import requests
import tempfile
import torch
import numpy as np

model_id = "csaybar/ldm-super-resolution-4x-cloudsen12"

# load model and scheduler
pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)

# load image
demo_file = "https://huggingface.co/csaybar/ldm-super-resolution-4x-cloudsen12/resolve/main/demo.pt"

# create temp file but keep it open
file = tempfile.NamedTemporaryFile()

with requests.get(demo_file, stream=True) as r:
    with open(file.name, 'wb') as f:
        for chunk in r.iter_content(chunk_size=8192):
            f.write(chunk)
    
    
torch_image = torch.load(file.name, map_location=torch.device('cpu'))

upscaled_image = pipeline(torch_image[0:1], num_inference_steps=100, output_type="numpy")
hrrgb = upscaled_image[0][0, :, :]
hrrgb = hrrgb[:, :, [0, 1, 2]]

lrrgb = torch_image[0, [0, 1, 2], :, :]
lrrgb = lrrgb.permute(1, 2, 0)
lrrgb = (lrrgb + 1)/2


import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, 2, figsize=(20, 10))
ax[0].imshow(lrrgb)
ax[1].imshow(hrrgb)
plt.show()