Update app.py

app.py

@@ -1,36 +1,28 @@
-import spaces
+import spaces 
 import os
 import sys
 import subprocess
 import tempfile
 import warnings
 warnings.filterwarnings('ignore')
-
 # ====================== DEPENDENCY SETUP ======================
 def setup():
-    """Fixed setup: Clone repo with submodules + install flash-attn properly + ensure kimia_infer import"""
+    """Fixed setup: Clone repo with submodules + install flash-attn properly"""
     print("🔧 Setting up dependencies...")
    
-    # 0. Install torch early (required for builds like flash-attn and pyproject.toml)
-    print("📦 Installing base torch...")
+    # 0. Install base torch with compatible versions and CUDA
+    print("📦 Installing base torch, torchaudio, torchvision...")
     try:
         subprocess.run([
             sys.executable, '-m', 'pip', 'install', '-q',
-            'torch==2.6.0', 'torchaudio==2.6.0'
+            'torch==2.6.0', 'torchaudio==2.6.0', 'torchvision==0.21.0',
+            '--index-url', 'https://download.pytorch.org/whl/cu126'
         ], check=True, stdout=sys.stdout, stderr=sys.stderr)
-        print("✅ torch installed")
+        print("✅ torch ecosystem installed")
     except subprocess.CalledProcessError as e:
         print(f"⚠️ torch install failed: {e}")
    
-    # 1. Flash-Attn
-    print("⚡ Installing flash-attn...")
-    try:
-        subprocess.run([
-            sys.executable, '-m', 'pip', 'install', '-q',
-            'flash-attn==2.7.4.post1', '--no-build-isolation'
-        ], check=True, stdout=sys.stdout, stderr=sys.stderr)
-    except subprocess.CalledProcessError as e:
-        print(f"⚠️ flash-attn failed: {e}")
+    # 1. Flash-Attn (install later after requirements)
    
     # 2. Clone Kimi-Audio with submodules
     repo_dir = "/tmp/Kimi-Audio"
@@ -54,7 +46,18 @@ def setup():
     except Exception as e:
         print(f"⚠️ requirements install failed: {e}")
    
-    # 4. Optional: Try editable install (but not critical now)
+    # 4. Force rebuild flash-attn from source to match torch
+    print("⚡ Forcing flash-attn build from source...")
+    try:
+        subprocess.run([
+            sys.executable, '-m', 'pip', 'install', '-q', 'flash-attn',
+            '--no-binary', 'flash-attn', '--force-reinstall', '--no-build-isolation'
+        ], check=True, stdout=sys.stdout, stderr=sys.stderr)
+        print("✅ flash-attn rebuilt")
+    except Exception as e:
+        print(f"⚠️ flash-attn rebuild failed: {e}")
+   
+    # 5. Optional: Try editable install
     print("🎵 Trying to install kimia_infer editable...")
     try:
         subprocess.run([
@@ -63,11 +66,11 @@ def setup():
     except Exception as e:
         print(f"⚠️ Editable install failed (ignoring, using path fallback): {e}")
    
-    # 5. Fallback: Add repo to sys.path for direct import
+    # 6. Fallback: Add repo to sys.path for direct import
     sys.path.insert(0, repo_dir)
     print(f"✅ Added {repo_dir} to sys.path: {sys.path[:2]}")  # Debug
    
-    # 6. Install other deps (overlaps with requirements.txt, but ensures HF-specific; remove torch since installed early)
+    # 7. Install other deps
     print("📚 Installing additional deps...")
     subprocess.run([
         sys.executable, '-m', 'pip', 'install', '-q',
@@ -75,7 +78,7 @@ def setup():
         'soundfile', 'gradio', 'spaces', 'pillow', 'numpy', 'scipy'
     ], check=True, stdout=sys.stdout, stderr=sys.stderr)
    
-    # 7. Early import test
+    # 8. Early import test
     try:
         from kimia_infer.api.kimia import KimiAudio
         print("✅ Early import test: kimia_infer SUCCESS")
@@ -83,10 +86,8 @@ def setup():
         print(f"❌ Early import test failed: {e}")
    
     print("✅ Setup completed!")
-
 # Run setup before any imports
 setup()
-
 # ====================== IMPORTS ======================
 import torch
 import gradio as gr
@@ -95,7 +96,6 @@ from huggingface_hub import snapshot_download
 import soundfile as sf
 from PIL import Image
 import numpy as np
-
 # Now safe to import kimia
 try:
     from kimia_infer.api.kimia import KimiAudio
@@ -105,7 +105,6 @@ except Exception as e:
     print(f"⚠️ KimiAudio import failed: {e}")
     KIMI_AUDIO_AVAILABLE = False
     KimiAudio = None
-
 # Try to import transformers for Kimi-VL
 try:
     from transformers import AutoProcessor, AutoModelForVision2Seq
@@ -115,11 +114,9 @@ except ImportError:
     KIMI_VL_AVAILABLE = False
     AutoProcessor = None
     AutoModelForVision2Seq = None
-
 print(f"CUDA available: {torch.cuda.is_available()}")
 if torch.cuda.is_available():
     print(f"GPU: {torch.cuda.get_device_name(0)}")
-
 # ====================== MODEL LOADING ======================
 class ModelManager:
     def __init__(self):
@@ -128,13 +125,13 @@ class ModelManager:
         self.vl_model = None
         self.vl_processor = None
         self.vl_device = None
-        
+       
     @spaces.GPU(duration=120)
     def load_audio_model(self):
         """Load Kimi-Audio with ZeroGPU"""
         if not KIMI_AUDIO_AVAILABLE:
             return "❌ kimia_infer not available"
-            
+           
         try:
             print("⬇️ Downloading Kimi-Audio-7B...")
             model_path = snapshot_download(
@@ -143,69 +140,67 @@ class ModelManager:
                 local_dir_use_symlinks=False,
                 resume_download=True
             )
-            
+           
             print(f"🚀 Loading Audio model...")
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            
+           
             model = KimiAudio(
                 model_path=model_path,
                 load_detokenizer=True
             )
             model = model.to(device)
-            
+           
             self.audio_model = model
             self.audio_device = device
             return f"✅ Audio model loaded on {device}"
         except Exception as e:
             return f"❌ Audio load failed: {str(e)}"
-    
+   
     @spaces.GPU(duration=180)
     def load_vl_model(self):
         """Load Kimi-VL with ZeroGPU"""
         if not KIMI_VL_AVAILABLE:
             return "❌ Transformers not available"
-            
+           
         try:
             print("⬇️ Downloading Kimi-VL-A3B...")
             model_id = "moonshotai/Kimi-VL-A3B-Thinking-2506"
-            
+           
             processor = AutoProcessor.from_pretrained(
-                model_id, 
+                model_id,
                 trust_remote_code=True
             )
-            
+           
             model = AutoModelForVision2Seq.from_pretrained(
                 model_id,
                 torch_dtype=torch.float16,
                 device_map="auto",
                 trust_remote_code=True
             )
-            
+           
             self.vl_processor = processor
             self.vl_model = model
             self.vl_device = next(model.parameters()).device
             return f"✅ VL model loaded on {self.vl_device}"
         except Exception as e:
             return f"❌ VL load failed: {str(e)}"
-
 # Global model manager
 manager = ModelManager()
-
 # ====================== INFERENCE FUNCTIONS ======================
 def generate_audio_response(audio_path: str, prompt: str):
     """Kimi-Audio inference"""
     if not manager.audio_model:
         return "Model not loaded. Click 'Load Audio Model' first.", None
-    
+   
     if not audio_path:
         return "Please upload audio.", None
-    
+   
     try:
         messages = [
             {"role": "user", "message_type": "text", "content": prompt or "Respond naturally."},
             {"role": "user", "message_type": "audio", "content": audio_path},
         ]
-        
+       
         sampling_params = {
             "audio_temperature": 0.8,
             "audio_top_k": 10,
@@ -216,40 +211,39 @@ def generate_audio_response(audio_path: str, prompt: str):
             "text_repetition_penalty": 1.0,
             "text_repetition_window_size": 16,
         }
-        
+       
         wav_output, text_output = manager.audio_model.generate(
             messages, **sampling_params, output_type="both"
         )
-        
+       
         # Save audio
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
             output_path = f.name
             if isinstance(wav_output, torch.Tensor):
                 wav_output = wav_output.detach().cpu().view(-1).numpy()
             sf.write(output_path, wav_output, 24000)
-        
+       
         return text_output, output_path
     except Exception as e:
         return f"Error: {str(e)}", None
-
 def generate_vl_response(image, text: str):
     """Kimi-VL inference"""
     if not manager.vl_model:
         return "Model not loaded. Click 'Load VL Model' first."
-    
+   
     if image is None:
         return "Please upload an image."
-    
+   
     try:
         # Format prompt for Kimi-VL
         prompt = f"<|im_start|>user\n<image>\n{text}<|im_end|>\n<|im_start|>assistant\n"
-        
+       
         inputs = manager.vl_processor(
             text=text,
             images=image,
             return_tensors="pt"
         ).to(manager.vl_device)
-        
+       
         outputs = manager.vl_model.generate(
             **inputs,
             max_new_tokens=512,
@@ -257,50 +251,48 @@ def generate_vl_response(image, text: str):
             temperature=0.7,
             top_p=0.9
         )
-        
+       
         response = manager.vl_processor.decode(outputs[0], skip_special_tokens=True)
         # Clean up the response (remove the prompt part)
         if "assistant" in response:
             response = response.split("assistant")[-1].strip()
-        
+       
         return response
     except Exception as e:
         return f"Error: {str(e)}"
-
 def chain_vl_to_audio(image, vl_prompt: str, audio_prompt: str):
     """Pipeline: Image → Kimi-VL description → Kimi-Audio narration"""
     if not manager.vl_model or not manager.audio_model:
         return "Both models must be loaded first.", None, None
-    
+   
     # Step 1: VL generates description
     description = generate_vl_response(image, vl_prompt)
-    
+   
     # Step 2: Audio generates speech from description
     # Create a dummy audio input for the text-to-speech mode if supported
     # Or use the description as text input to audio model
     text_out, audio_out = generate_audio_response(None, f"Narrate this: {description}")
-    
+   
     return description, text_out, audio_out
-
 # ====================== GRADIO UI ======================
 with gr.Blocks(title="Kimi Multimodal Lab • ZeroGPU", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🎭🎵👁️ Kimi Multimodal Test Lab
     **Kimi-Audio-7B** (Voice) + **Kimi-VL-A3B** (Vision) on HuggingFace ZeroGPU
     """)
-    
+   
     with gr.Tab("🚀 Model Setup"):
         gr.Markdown("Load models first (takes 60-120s each on ZeroGPU)")
         with gr.Row():
             load_audio_btn = gr.Button("Load Kimi-Audio", variant="primary")
             load_vl_btn = gr.Button("Load Kimi-VL", variant="primary")
-        
+       
         audio_status = gr.Textbox(label="Audio Model Status", value="Not loaded")
         vl_status = gr.Textbox(label="VL Model Status", value="Not loaded")
-        
+       
         load_audio_btn.click(manager.load_audio_model, outputs=audio_status)
         load_vl_btn.click(manager.load_vl_model, outputs=vl_status)
-    
+   
     with gr.Tab("🎵 Kimi-Audio"):
         gr.Markdown("Voice conversation, ASR, audio Q&A")
         with gr.Row():
@@ -316,17 +308,17 @@ with gr.Blocks(title="Kimi Multimodal Lab • ZeroGPU", theme=gr.themes.Soft())
                     placeholder="E.g., 'What is being said?' or 'Summarize the meeting'"
                 )
                 audio_gen_btn = gr.Button("Generate Response", variant="primary")
-            
+           
             with gr.Column():
                 audio_text_out = gr.Textbox(label="Text Response", lines=4)
                 audio_out = gr.Audio(label="Kimi's Voice Response", type="filepath")
-        
+       
         audio_gen_btn.click(
             generate_audio_response,
             inputs=[audio_input, audio_text_prompt],
             outputs=[audio_text_out, audio_out]
         )
-    
+   
     with gr.Tab("👁️ Kimi-VL"):
         gr.Markdown("Visual question answering, image description, visual comedy")
         with gr.Row():
@@ -338,16 +330,16 @@ with gr.Blocks(title="Kimi Multimodal Lab • ZeroGPU", theme=gr.themes.Soft())
                     placeholder="E.g., 'What do you see?' or 'Roast this outfit'"
                 )
                 vl_gen_btn = gr.Button("Analyze Image", variant="primary")
-            
+           
             with gr.Column():
                 vl_output = gr.Textbox(label="Visual Analysis", lines=8)
-        
+       
         vl_gen_btn.click(
             generate_vl_response,
             inputs=[image_input, vl_text_prompt],
             outputs=vl_output
         )
-    
+   
     with gr.Tab("🎭 Combined Pipeline"):
         gr.Markdown("Chain: Image → Description → Voice Narration")
         with gr.Row():
@@ -362,18 +354,18 @@ with gr.Blocks(title="Kimi Multimodal Lab • ZeroGPU", theme=gr.themes.Soft())
                     label="Voice Style Prompt"
                 )
                 chain_btn = gr.Button("Run Full Pipeline", variant="primary")
-            
+           
             with gr.Column():
                 chain_desc = gr.Textbox(label="Generated Description")
                 chain_text = gr.Textbox(label="Audio Text")
                 chain_audio = gr.Audio(label="Narrated Audio")
-        
+       
         chain_btn.click(
             chain_vl_to_audio,
             inputs=[chain_image, chain_vl_prompt, chain_audio_prompt],
             outputs=[chain_desc, chain_text, chain_audio]
         )
-    
+   
     gr.Markdown("---")
     gr.Markdown("""
     **Notes:**
@@ -381,14 +373,10 @@ with gr.Blocks(title="Kimi Multimodal Lab • ZeroGPU", theme=gr.themes.Soft())
     - ZeroGPU provides A100/L4 GPUs - cold start ~60-120s per model
     - Keep `max_size=1` in queue to prevent OOM with two large models
     """)
-
-
 import asyncio
 import warnings
-
 # Suppress the event loop cleanup error
 warnings.filterwarnings("ignore", category=ResourceWarning)
-
 # Fix for asyncio cleanup on exit
 def silence_event_loop_closed(func):
     def wrapper(*args, **kwargs):
@@ -400,17 +388,15 @@ def silence_event_loop_closed(func):
             else:
                 raise
     return wrapper
-
 # Patch the event loop to prevent the error
 asyncio.base_events.BaseEventLoop.__del__ = silence_event_loop_closed(
     asyncio.base_events.BaseEventLoop.__del__
 )
-
 # Disable SSR (experimental mode causing the issue)
 demo.queue(max_size=1)
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
-    share=False,  # Set to True if you need a public gradio.live link
-    ssr_mode=False  # <-- DISABLES the experimental SSR causing the error
+    share=False, # Set to True if you need a public gradio.live link
+    ssr_mode=False # <-- DISABLES the experimental SSR causing the error
 )