fix output

app.py

@@ -173,8 +173,13 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
         image_tensor = get_value_at_index(vaedecode_79, 0)
         # Convert from [0,1] to [0,255] range and ensure it's in the right format
         image_np = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
-        # Reshape from (batch, channel, height, width) to (height, width, channel)
-        image_np = np.transpose(image_np[0], (1, 2, 0))
+        
+        # Ensure we have the right shape (height, width, channels)
+        if len(image_np.shape) == 4:  # (batch, channel, height, width)
+            image_np = image_np[0]  # Remove batch dimension
+        if len(image_np.shape) == 3:  # (channel, height, width)
+            image_np = np.transpose(image_np, (1, 2, 0))  # Convert to (height, width, channel)
+        
         # Convert to PIL Image
         image = Image.fromarray(image_np)
         return image