Update app.py

app.py

@@ -221,10 +221,16 @@ def infer(style_description, ref_style_file, caption):
         ],
         dim=0)
 
-    # Save the sampled image to a file
-    sampled_image = T.ToPILImage()(sampled.squeeze(0))  # Convert tensor to PIL image
-    sampled_image.save(output_file)  # Save the image
-
+    # Remove batch dimension if it exists
+    if sampled.dim() == 4 and sampled.size(0) == 1:
+        sampled = sampled.squeeze(0)
+
+    # Ensure the tensor is in [C, H, W] format
+    if sampled.dim() == 3:
+        sampled_image = T.ToPILImage()(sampled)  # Convert tensor to PIL image
+        sampled_image.save(output_file)  # Save the image as a PNG
+    else:
+        raise ValueError(f"Expected tensor of shape [C, H, W] but got {sampled.shape}")
     clear_gpu_cache()  # Clear cache after inference
 
     return output_file  # Return the path to the saved image