Update app.py to use the new SonicDiffusion controller

app.py

@@ -15,7 +15,7 @@ os.makedirs("outputs", exist_ok=True)
 # Import required packages
 import gradio as gr
 
-# Try importing key packages
+# Try importing key packages and print versions
 packages = {
     "torch": None,
     "transformers": None,
@@ -26,54 +26,139 @@ packages = {
 for package in packages.keys():
     try:
         module = __import__(package)
-        packages[package] = module.__version__
-        print(f"{package} version: {module.__version__}")
+        try:
+            packages[package] = module.__version__
+            print(f"{package} version: {module.__version__}")
+        except AttributeError:
+            packages[package] = "Installed (version unknown)"
+            print(f"{package} installed (version unknown)")
     except ImportError as e:
         print(f"{package} import error: {e}")
 
 # Import our controller
-from controller import SimpleSonicDiffusionController
+from controller import SonicDiffusionController
 
 # Initialize controller
-controller = SimpleSonicDiffusionController()
+controller = SonicDiffusionController()
 
 # Create the Gradio interface
-with gr.Blocks(title="SonicDiffusion - Progressive Setup") as demo:
-    gr.Markdown("# SonicDiffusion - Building Up")
+with gr.Blocks(title="SonicDiffusion") as demo:
+    gr.Markdown("# SonicDiffusion - Audio-to-Image Generation")
+    gr.Markdown("Generate images conditioned on audio inputs using Stable Diffusion")
     
-    status_output = gr.Textbox(label="Status", value="System initialized. Click 'Check System' to verify setup.")
+    status_output = gr.Textbox(label="Status", value="System initialized. Check dependencies and download assets first.")
     
-    with gr.Tab("System Check"):
-        check_btn = gr.Button("Check System")
-        
-        def check_system():
-            status = ["Package Status:"]
-            
-            # Check package availability
-            for package, version in packages.items():
-                status.append(f"{package}: {version if version else 'Not Available'}")
+    with gr.Tab("1. Setup"):
+        with gr.Row():
+            with gr.Column():
+                check_deps_btn = gr.Button("Check Dependencies")
+                
+                def format_deps(deps):
+                    return "\n".join([f"{pkg}: {vers}" for pkg, vers in deps.items()])
+                
+                check_deps_btn.click(
+                    fn=lambda: format_deps(controller.check_dependencies()),
+                    outputs=status_output
+                )
             
-            # Check directories
-            status.append("\nDirectory Status:")
-            asset_status = controller.get_asset_status()
-            for dir_name, dir_status in asset_status.items():
-                if dir_name in ["assets", "ckpts", "outputs"]:
-                    status.append(f"Directory '{dir_name}': {dir_status}")
+            with gr.Column():
+                check_assets_btn = gr.Button("Check Assets")
                 
-            return "\n".join(status)
-        
-        check_btn.click(fn=check_system, outputs=status_output)
-    
-    with gr.Tab("Model"):
-        load_model_btn = gr.Button("Load Model")
-        load_model_btn.click(fn=controller.load_model, outputs=status_output)
+                def format_assets(assets):
+                    return "\n".join([f"{path}: {'✓' if exists else '✗'}" for path, exists in assets.items()])
+                
+                check_assets_btn.click(
+                    fn=lambda: format_assets(controller.check_assets()),
+                    outputs=status_output
+                )
+                
+        download_assets_btn = gr.Button("Download Required Assets")
+        download_assets_btn.click(fn=controller.download_assets, outputs=status_output)
     
-    with gr.Tab("Generate"):
-        text_input = gr.Textbox(label="Enter a prompt", value="Hello, SonicDiffusion!")
-        gen_btn = gr.Button("Process Text")
-        gen_output = gr.Textbox(label="Output")
+    with gr.Tab("2. Generate"):
+        with gr.Row():
+            with gr.Column():
+                model_dropdown = gr.Dropdown(
+                    label="Select Model",
+                    choices=["Landscape Model", "Greatest Hits Model"],
+                    value="Landscape Model"
+                )
+                
+                load_model_btn = gr.Button("Load Selected Model")
+                load_model_btn.click(
+                    fn=controller.load_model,
+                    inputs=[model_dropdown],
+                    outputs=status_output
+                )
+                
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter a descriptive prompt...",
+                    value="a high quality photograph of a fantasy landscape"
+                )
+                
+                audio_input = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    sources=["upload"]
+                )
+                
+                with gr.Row():
+                    cfg_scale = gr.Slider(
+                        label="CFG Scale",
+                        minimum=1.0,
+                        maximum=20.0,
+                        value=7.5,
+                        step=0.5
+                    )
+                    
+                    steps = gr.Slider(
+                        label="Steps",
+                        minimum=20,
+                        maximum=100,
+                        value=50,
+                        step=5
+                    )
+                
+                generate_btn = gr.Button("Generate", variant="primary")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Generated Image")
+                
+                generate_btn.click(
+                    fn=controller.generate,
+                    inputs=[prompt_input, audio_input, cfg_scale, steps],
+                    outputs=output_image
+                )
         
-        gen_btn.click(fn=controller.generate, inputs=[text_input], outputs=gen_output)
+        with gr.Row():
+            gr.Markdown("### Example Audio Files")
+            examples = [
+                ["a serene landscape with mountains and a lake", "assets/fire_crackling.wav", 7.5, 50],
+                ["a mysterious forest at night with glowing elements", "assets/plastic_bag.wav", 7.5, 50]
+            ]
+            gr.Examples(
+                examples=examples,
+                inputs=[prompt_input, audio_input, cfg_scale, steps],
+                outputs=output_image,
+                fn=controller.generate
+            )
 
 if __name__ == "__main__":
+    # Attempt to download example audio files if they don't exist
+    if not os.path.exists("assets/fire_crackling.wav") or not os.path.exists("assets/plastic_bag.wav"):
+        try:
+            from download_assets import download_gdrive_file
+            
+            assets = {
+                "assets/fire_crackling.wav": "1vOAZcbkpo_hre2g26n--lUXdwbTQp22k",
+                "assets/plastic_bag.wav": "15igeDor7a47a-oluSCfO6GeUvFVl2ttb"
+            }
+            
+            for path, file_id in assets.items():
+                if not os.path.exists(path):
+                    download_gdrive_file(file_id, path)
+        except Exception as e:
+            print(f"Error downloading example audio files: {e}")
+    
     demo.launch()