Update app.py with better error handling and audio selection

app.py

@@ -7,11 +7,51 @@ print(f"Python version: {sys.version}")
 print(f"Working directory: {os.getcwd()}")
 print(f"Directory contents: {os.listdir('.')}")
 
+# Import key packages and print versions
+packages = {
+    "torch": None,
+    "transformers": None,
+    "diffusers": None,
+    "accelerate": None,
+    "einops": None,
+    "omegaconf": None,
+    "librosa": None 
+}
+
+for package in packages.keys():
+    try:
+        module = __import__(package)
+        try:
+            packages[package] = module.__version__
+            print(f"{package}: {module.__version__}")
+        except AttributeError:
+            packages[package] = "Installed (version unknown)"
+            print(f"{package}: Installed (version unknown)")
+    except ImportError as e:
+        print(f"{package}: Not installed")
+        packages[package] = "Not installed"
+
 # Create necessary directories
 os.makedirs("assets", exist_ok=True)
 os.makedirs("ckpts", exist_ok=True)
 os.makedirs("outputs", exist_ok=True)
 
+# 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}")
+
 # Import required packages
 import gradio as gr
 
@@ -26,7 +66,7 @@ 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. Check dependencies and download assets first.")
+    status_output = gr.Textbox(label="Status", value="System initialized. Start by checking dependencies and downloading assets.", lines=8)
     
     with gr.Tab("1. Setup"):
         with gr.Row():
@@ -54,7 +94,7 @@ with gr.Blocks(title="SonicDiffusion") as demo:
                     outputs=status_output
                 )
                 
-        download_assets_btn = gr.Button("Download Required Assets")
+        download_assets_btn = gr.Button("Download Required Assets", variant="primary")
         download_assets_btn.click(fn=controller.download_assets, outputs=status_output)
     
     with gr.Tab("2. Generate"):
@@ -66,7 +106,7 @@ with gr.Blocks(title="SonicDiffusion") as demo:
                     value="Landscape Model"
                 )
                 
-                load_model_btn = gr.Button("Load Selected Model")
+                load_model_btn = gr.Button("Load Selected Model", variant="primary")
                 load_model_btn.click(
                     fn=controller.load_model,
                     inputs=[model_dropdown],
@@ -79,10 +119,27 @@ with gr.Blocks(title="SonicDiffusion") as demo:
                     value="a high quality photograph of a fantasy landscape"
                 )
                 
+                # Use the preloaded example audio files
+                audio_dropdown = gr.Dropdown(
+                    label="Select Example Audio",
+                    choices=["assets/fire_crackling.wav", "assets/plastic_bag.wav", "Upload my own audio"],
+                    value="assets/fire_crackling.wav"
+                )
+                
                 audio_input = gr.Audio(
-                    label="Upload Audio",
+                    label="Or Upload Your Own Audio",
                     type="filepath",
-                    sources=["upload"]
+                    sources=["upload", "microphone"],
+                    visible=False
+                )
+                
+                def toggle_audio_upload(choice):
+                    return gr.update(visible=(choice == "Upload my own audio"))
+                
+                audio_dropdown.change(
+                    fn=toggle_audio_upload,
+                    inputs=[audio_dropdown],
+                    outputs=[audio_input]
                 )
                 
                 with gr.Row():
@@ -102,42 +159,26 @@ with gr.Blocks(title="SonicDiffusion") as demo:
                         step=5
                     )
                 
-                generate_btn = gr.Button("Generate", variant="primary")
+                def prepare_audio_path(dropdown_choice, uploaded_audio):
+                    if dropdown_choice == "Upload my own audio":
+                        return uploaded_audio
+                    else:
+                        return dropdown_choice
+                
+                generate_btn = gr.Button("Generate Image", variant="primary")
             
             with gr.Column():
-                output_image = gr.Image(label="Generated Image")
+                output_image = gr.Image(label="Generated Image", type="pil")
+                output_status = gr.Textbox(label="Generation Status", value="")
                 
                 generate_btn.click(
-                    fn=controller.generate,
-                    inputs=[prompt_input, audio_input, cfg_scale, steps],
-                    outputs=output_image
+                    fn=lambda p, d, u, c, s: (
+                        controller.generate(p, prepare_audio_path(d, u), c, s),
+                        "Generation completed"
+                    ),
+                    inputs=[prompt_input, audio_dropdown, audio_input, cfg_scale, steps],
+                    outputs=[output_image, output_status]
                 )
-        
-        with gr.Row():
-            gr.Markdown("### Example Audio Files")
-            gr.Markdown("""
-            Example audio files will be downloaded to the 'assets' folder when you click 'Download Required Assets'.
-            - fire_crackling.wav - Natural fire sounds
-            - plastic_bag.wav - Crinkling plastic sounds
-            
-            Use these with prompts like "a serene landscape with mountains" or "a mysterious forest at night".
-            """)
 
 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()