fix output

app.py

@@ -5,6 +5,7 @@ from typing import Sequence, Mapping, Any, Union
 import torch
 import gradio as gr
 from PIL import Image
+import numpy as np
 from huggingface_hub import hf_hub_download
 import spaces
 from comfy import model_management
@@ -168,8 +169,13 @@ def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
             vae=get_value_at_index(vaeloader_80, 0),
         )
 
-        # Instead of saving to file, return the image directly
-        return get_value_at_index(vaedecode_79, 0)
+        # Convert tensor to PIL Image
+        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)
+        # Convert to PIL Image
+        image = Image.fromarray(image_np)
+        return image
 
 # Create Gradio interface
 with gr.Blocks() as app: