Update app.py

app.py

@@ -37,8 +37,13 @@ class StableBuddyApp:
             else:
                 image = self.pipe(prompt, guidance_scale=8.5).images[0]
             
+            # Ensure the directory exists
+            output_dir = 'data'
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)
+
             # Save the generated image temporarily
-            image_path = 'data/generated_image.png'
+            image_path = os.path.join(output_dir, 'generated_image.png')
             image.save(image_path)
             
             return image_path  # Return the image path for Gradio to display