Update app.py

app.py

@@ -28,8 +28,27 @@ model.load_state_dict(ckpt)
 diffusion = Diffusion_cond(img_size=256, device=device)
 model.eval()
 
-transform_hmi = transforms.Compose([
+from torchvision import transforms
+
+# Define a custom transform to clamp data
+class ClampTransform(object):
+    def __init__(self, min_value=-250, max_value=250):
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def __call__(self, tensor):
+        return torch.clamp(tensor, self.min_value, self.max_value)
+
+transform_hmi_jp2 = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Resize((256, 256)),
+    transforms.RandomVerticalFlip(p=1.0),
+    transforms.Normalize(mean=(0.5,), std=(0.5,))
+])
+
+transform_hmi_fits = transforms.Compose([
     transforms.ToTensor(),
+    ClampTransform(-250, 250),
     transforms.Resize((256, 256)),
     transforms.RandomVerticalFlip(p=1.0),
     transforms.Normalize(mean=(0.5,), std=(0.5,))
@@ -39,12 +58,12 @@ def generate_image(seed_image):
     _, file_ext = os.path.splitext(seed_image)
     
     if file_ext.lower() == '.jp2':
-        input_img = Image.open(seed_image)
+        input_img = Image.transform_hmi_jp2(seed_image)
         input_img_pil = transform_hmi(input_img).reshape(1, 1, 256, 256).to(device)
     elif file_ext.lower() == '.fits':
         with fits.open(seed_image) as hdul:
             data = hdul[0].data
-            input_img_pil = transform_hmi(data).reshape(1, 1, 256, 256).to(device)
+            input_img_pil = transform_hmi_fits(data).reshape(1, 1, 256, 256).to(device)
     else:
         print(f'Format {file_ext.lower()} not supported')