import gradio as gr
import tempfile
import os
from SDXLImageGenerator import SDXLImageGenerator  # Import your existing class
import sys
from Image3DProcessor import Image3DProcessor  # Import your 3D processing class

# Ensure setup.sh runs and paths are appended
os.system('bash setup.sh')
sys.path.append('./splatter-image')
sys.path.append('./diff-gaussian-rasterization')

class ControlNetProcessor:
    def controlnet_image(self, image):
        # Placeholder for ControlNet processing (e.g., returning a processed image or placeholder text)
        return image  # Returning the image for further processing

class VideoGenerator:
    def __init__(self, model_cfg_path, model_repo_id, model_filename):
        # Initialize the Image3DProcessor
        self.processor = Image3DProcessor(model_cfg_path, model_repo_id, model_filename)

    def generate_3d_video(self, image):
        # Process the image and create a 3D video
        processed_image = self.processor.preprocess(image, preprocess_background=False)
        mesh_path, video_path = self.processor.reconstruct_and_export(processed_image)

        return video_path

class GradioApp:
    def __init__(self):
        self.sdxl_generator = SDXLImageGenerator()  # Use your existing class
        self.controlnet_processor = ControlNetProcessor()
        # Initialize VideoGenerator with required paths and details
        self.video_generator = VideoGenerator(
            model_cfg_path="path/to/gradio_config.yaml",
            model_repo_id="szymanowiczs/splatter-image-multi-category-v1",
            model_filename="model_latest.pth"
        )

    def full_pipeline(self, prompt):
        initial_image = self.sdxl_generator.generate_images([prompt])[0]
        controlled_image = self.controlnet_processor.controlnet_image(initial_image)
        video_path = self.video_generator.generate_3d_video(controlled_image)
        return initial_image, controlled_image, video_path

    def launch(self):
        interface = gr.Interface(
            fn=self.full_pipeline,
            inputs=gr.Textbox(label="Input Prompt"),
            outputs=[
                gr.Image(label="Generated Image"),
                gr.Image(label="ControlNet Processed Image"),
                gr.Video(label="3D Model Video")
            ],
            title="SDXL to ControlNet to 3D Pipeline",
            description="Generate an image using SDXL, refine it with ControlNet, and generate a 3D video output."
        )
        interface.launch(share=True)  # Added `share=True` for public link

if __name__ == "__main__":
    app = GradioApp()
    app.launch()