fixed impot

app.py

@@ -156,7 +156,7 @@ demo = gr.Interface(
         gr.Image(type="filepath", label="Input"),
         gr.Radio(['aligned', 'unaligned'], type="value", value='unaligned', label='Image Alignment'),
         gr.Number(label="Rescaling factor", value=2),
-        gr.Number(label="Number of flow steps", value=25),
+        gr.Number(label="Number of flow steps (a higher value leads to better image quality at the expense of runtime)", value=25),
     ], [
         gr.Image(type="numpy", label="Output (The whole image)"),
         gr.File(label="Download the output image")

lightning_models/mmse_rectified_flow.py

@@ -8,7 +8,8 @@ from torch.nn.functional import mse_loss
 from torch.nn.functional import sigmoid
 from torch.optim import AdamW
 from torch_ema import ExponentialMovingAverage as EMA
-from torchmetrics.image import FrechetInceptionDistance, InceptionScore
+from torchmetrics.image.fid import FrechetInceptionDistance
+from torchmetrics.image.inception import InceptionScore
 from torchvision.transforms.functional import to_pil_image
 from torchvision.utils import save_image