Code now works

SDXLImageGenerator.py

@@ -22,7 +22,7 @@ class SDXLImageGenerator:
         )
         self.pipe.to(self.device)
 
-    def generate_images(self, prompts):
+    def generate_image(self, prompts):
         start_time = time.time()
         
         # Generate images in a batch

app.py

@@ -6,6 +6,8 @@ import sys
 from Image3DProcessor import Image3DProcessor  # Import your 3D processing class
 from PIL import Image
 import io
+from io import BytesIO
+import numpy as np
 
 class VideoGenerator:
     def __init__(self, model_cfg_path, model_repo_id, model_filename):
@@ -13,15 +15,21 @@ class VideoGenerator:
         self.processor = Image3DProcessor(model_cfg_path, model_repo_id, model_filename)
 
     def generate_3d_video(self, image):
+        # Ensure the image is a PIL Image object
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
         # Preprocess the image first
         processed_image = self.processor.preprocess(image)
         # Then pass it to reconstruct_and_export
         video_data = self.processor.reconstruct_and_export(processed_image)
-        return video_data
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as video_file:
+            video_file.write(video_data)
+            video_path = video_file.name
+        return video_path
 
 class GradioApp:
     def __init__(self):
-        self.sdxl_generator = SDXLImageGenerator()  # Use your existing class
+        self.sdxl_generator = SDXLImageGenerator()
         # Initialize VideoGenerator with required paths and details
         self.video_generator = VideoGenerator(
             model_cfg_path="/home/user/app/splatter-image/gradio_config.yaml",
@@ -29,34 +37,53 @@ class GradioApp:
             model_filename="model_latest.pth"
         )
 
-    def full_pipeline(self, prompt):
-        # Generate the initial image using SDXLImageGenerator
-        initial_image = self.sdxl_generator.generate_images([prompt])[0]
-
-        # Generate a 3D video using the image
-        video_data = self.video_generator.generate_3d_video(initial_image)
-
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as video_file:
-            video_file.write(video_data)
-            video_path = video_file.name
-
-        # Convert bytes to a PIL Image for further processing and display
-        initial_image = Image.open(io.BytesIO(initial_image))
-        
-        return initial_image, video_path
-
     def launch(self):
         with gr.Blocks() as interface:
+            # Input for the prompt at the top
             prompt_input = gr.Textbox(label="Input Prompt", elem_id="input_textbox")
-            generate_button = gr.Button("Generate")
+
+            # Button for generating the 3D object
+            generate_3d_object = gr.Button("Generate 3D object")
+
+            # Outputs: image on the bottom left, video on the bottom right
             with gr.Row():
-                image_output = gr.Image(label="Generated Image", elem_id="generated_image")
-                video_output = gr.Video(label="3D Model Video", elem_id="model_video")
+                with gr.Column():
+                    image_output = gr.Image(label="Generated Image", elem_id="generated_image")
+                with gr.Column():
+                    video_output = gr.Video(label="3D Model Video", elem_id="model_video")
+
+            # Generate the image first
+            def generate_image_and_display(prompt):
+                # Generate the image from the prompt
+                image_data = self.sdxl_generator.generate_image([prompt])[0]
+                return Image.open(BytesIO(image_data))
+
+            # Generate the 3D after the image is ready
+            def generate_3D_from_image(image):
+                # Ensure the image is a PIL Image object
+                if isinstance(image, np.ndarray):
+                    image = Image.fromarray(image)
+                # Generate the 3D from the generated image
+                return self.video_generator.generate_3d_video(image)
+
+            # First click generates the image
+            generate_3d_object.click(
+                fn=generate_image_and_display,
+                inputs=prompt_input,
+                outputs=image_output,
+                queue=True
+            )
+
+            # Once the image is ready, generate the video
+            image_output.change(
+                fn=generate_3D_from_image,
+                inputs=image_output,
+                outputs=video_output,
+                queue=True
+            )
 
-            generate_button.click(fn=self.full_pipeline, inputs=prompt_input, outputs=[image_output, video_output])
-        
         interface.launch(share=True)
 
 if __name__ == "__main__":
     app = GradioApp()
-    app.launch()
+    app.launch()