Update app.py

app.py

@@ -1,168 +1,135 @@
 import gradio as gr
-from gliner import GLiNER
-from vllm import LLM, SamplingParams
-from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
-import json
 import torch
-import requests
-import threading
-from queue import Queue
 import logging
-import pynvml
+from transformers import AutoTokenizer, AutoModel
+from diffusers import DiffusionPipeline
+import soundfile as sf
+import numpy as np
 
-# Configure logging
-logging.basicConfig(level=logging.DEBUG)
+# Set up logging to debug startup issues
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Initialize NVML for GPU debugging
 try:
-    pynvml.nvmlInit()
-    device_count = pynvml.nvmlDeviceGetCount()
-    logger.info(f"NVML Initialized. GPU Count: {device_count}")
-    for i in range(device_count):
-        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-        name = pynvml.nvmlDeviceGetName(handle)
-        logger.info(f"GPU {i}: {name}")
-except pynvml.NVMLError as e:
-    logger.error(f"NVML Initialization Failed: {str(e)}")
-    raise RuntimeError("Cannot initialize NVML. Check NVIDIA drivers.")
-
-# Verify CUDA
-if not torch.cuda.is_available():
-    logger.error("CUDA not available")
-    raise RuntimeError("No GPU detected. Ensure H200 GPU is available.")
-logger.info(f"CUDA Version: {torch.version.cuda}")
-logger.info(f"GPU Detected: {torch.cuda.get_device_name(0)}")
-logger.info(f"Device Count: {torch.cuda.device_count()}")
-
-# Load legal corpus
-with open("legal_corpus.json", "r", encoding="utf-8") as f:
-    corpus = json.load(f)
-documents = [item["text"] for item in corpus]
-
-# Initialize sentence transformer (GPU)
-embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cuda")
-embeddings = embedder.encode(documents, convert_to_numpy=True)
-
-# Initialize FAISS-GPU
-dimension = embeddings.shape[1]
-index = faiss.IndexFlatL2(dimension)
-index.add(embeddings)
-
-# Initialize GLiNER (GPU)
-gliner_model = GLiNER.from_pretrained("NAMAA-Space/gliner_arabic-v2.1", load_tokenizer=True)
-gliner_model = gliner_model.cuda()
-
-# Initialize LLM (default to Qwen2-7B-Instruct-AWQ)
-use_qwq_32b = False  # Set to True if H200 detection is fixed
-model_name = "Qwen/Qwen2-7B-Instruct-AWQ" if not use_qwq_32b else "Qwen/QwQ-32B"
-try:
-    llm = LLM(
-        model=model_name,
-        quantization="awq",
-        max_model_len=4096,
-        gpu_memory_utilization=0.9,
-        device="cuda"
-    )
-    logger.info(f"Loaded LLM: {model_name}")
-except Exception as e:
-    logger.error(f"Failed to initialize LLM: {str(e)}")
-    raise
-
-sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
-
-def fetch_external_legal_data(query, queue):
-    """Fetch external legal data via HTTP request (mock API)."""
-    try:
-        response = requests.get(
-            "https://api.example.com/legal",
-            params={"query": query},
-            timeout=5
+    # Load text tokenizer and embedding model (umt5-base)
+    def load_text_processor():
+        logger.info("Loading text processor (umt5-base)...")
+        tokenizer = AutoTokenizer.from_pretrained("./umt5-base")
+        text_model = AutoModel.from_pretrained(
+            "./umt5-base",
+            use_safetensors=True,
+            torch_dtype=torch.float16,
+            device_map="auto"
         )
-        response.raise_for_status()
-        queue.put(response.json().get("text", "No external data found"))
-    except requests.RequestException:
-        queue.put("Failed to fetch external data")
-
-def run_ner(text, entity_types, queue):
-    """Run NER with gliner_arabic-v2.1."""
-    if not text or not entity_types:
-        queue.put([])
-        return
-    entity_list = [e.strip() for e in entity_types.split(",")]
-    entities = gliner_model.predict_entities(text, entity_list, threshold=0.5)
-    queue.put([{"text": e["text"], "label": e["label"], "score": round(e["score"], 2)} for e in entities])
-
-def retrieve_documents(query, k=2):
-    """Retrieve top-k documents using FAISS-GPU."""
-    query_embedding = embedder.encode([query], convert_to_numpy=True)
-    distances, indices = index.search(query_embedding, k)
-    return [documents[idx] for idx in indices[0]]
-
-def generate_legal_insight(text, entities, retrieved_docs, external_data):
-    """Generate insight with LLM using RAG."""
-    entity_str = ", ".join([f"{e['text']} ({e['label']})" for e in entities])
-    context = "\n".join(retrieved_docs) + "\nExternal Data: " + external_data
-    prompt = f"""You are a legal assistant for Arabic law. Using the following context, extracted entities, and external data, provide a concise legal insight.
-
-Context:
-{context}
-
-Entities:
-{entity_str}
-
-Input Text:
-{text}
-
-Insight:"""
-    outputs = llm.generate([prompt], sampling_params)
-    return outputs[0].outputs[0].text
-
-def main_interface(text, entity_types):
-    """Main Gradio interface with threading."""
-    ner_queue = Queue()
-    external_queue = Queue()
-
-    ner_thread = threading.Thread(target=run_ner, args=(text, entity_types, ner_queue))
-    external_thread = threading.Thread(target=fetch_external_legal_data, args=(text, external_queue))
-    
-    ner_thread.start()
-    external_thread.start()
-    
-    ner_thread.join()
-    external_thread.join()
-
-    ner_result = ner_queue.get()
-    external_data = external_queue.get()
-
-    retrieved_docs = retrieve_documents(text)
-    
-    insight = generate_legal_insight(text, ner_result, retrieved_docs, external_data)
-    
-    return ner_result, retrieved_docs, external_data, insight
-
-# Gradio interface
-with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-    gr.Markdown("# Arabic Legal Demo: NER & RAG with GLiNER and LLM")
-    with gr.Row():
-        text_input = gr.Textbox(label="Arabic Legal Text", lines=5, placeholder="Enter Arabic legal text...")
-        entity_types = gr.Textbox(
-            label="Entity Types (comma-separated)",
-            value="person,law,organization",
-            placeholder="e.g., person,law,organization"
+        logger.info("Text processor loaded successfully.")
+        return tokenizer, text_model
+
+    # Load the transformer backbone (phantomstep_transformer)
+    def load_transformer():
+        logger.info("Loading transformer (phantomstep_transformer)...")
+        transformer = DiffusionPipeline.from_pretrained(
+            "./phantomstep_transformer",
+            use_safetensors=True,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        logger.info("Transformer loaded successfully.")
+        return transformer
+
+    # Load the DCAE for audio encoding/decoding (phantomstep_dcae)
+    def load_dcae():
+        logger.info("Loading DCAE (phantomstep_dcae)...")
+        dcae = DiffusionPipeline.from_pretrained(
+            "./phantomstep_dcae",
+            use_safetensors=True,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        logger.info("DCAE loaded successfully.")
+        return dcae
+
+    # Load the vocoder for audio synthesis (phantomstep_vocoder)
+    def load_vocoder():
+        logger.info("Loading vocoder (phantomstep_vocoder)...")
+        vocoder = DiffusionPipeline.from_pretrained(
+            "./phantomstep_vocoder",
+            use_safetensors=True,
+            torch_dtype=torch.float16,
+            device_map="auto"
         )
-    submit_btn = gr.Button("Analyze")
-    ner_output = gr.JSON(label="Extracted Entities")
-    docs_output = gr.Textbox(label="Retrieved Legal Context")
-    external_output = gr.Textbox(label="External Legal Data")
-    insight_output = gr.Textbox(label="Legal Insight")
-    
-    submit_btn.click(
-        fn=main_interface,
-        inputs=[text_input, entity_types],
-        outputs=[ner_output, docs_output, external_output, insight_output]
-    )
+        logger.info("Vocoder loaded successfully.")
+        return vocoder
+
+    # Generate music from a text prompt
+    def generate_music(prompt, duration=20, seed=42):
+        logger.info(f"Generating music with prompt: {prompt}, duration: {duration}, seed: {seed}")
+        torch.manual_seed(seed)
+        
+        # Load all components
+        tokenizer, text_model = load_text_processor()
+        transformer = load_transformer()
+        dcae = load_dcae()
+        vocoder = load_vocoder()
+        
+        # Step 1: Process text prompt to embeddings
+        logger.info("Processing text prompt to embeddings...")
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        inputs = {k: v.to(text_model.device) for k, v in inputs.items()}
+        with torch.no_grad():
+            embeddings = text_model(**inputs).last_hidden_state.mean(dim=1)
+        
+        # Step 2: Pass embeddings through transformer
+        logger.info("Generating with transformer...")
+        transformer_output = transformer(
+            embeddings,
+            num_inference_steps=50,
+            audio_length_in_s=duration
+        ).audios[0]
+        
+        # Step 3: Decode audio features with DCAE
+        logger.info("Decoding with DCAE...")
+        dcae_output = dcae(
+            transformer_output,
+            num_inference_steps=50,
+            audio_length_in_s=duration
+        ).audios[0]
+        
+        # Step 4: Synthesize final audio with vocoder
+        logger.info("Synthesizing with vocoder...")
+        audio = vocoder(
+            dcae_output,
+            num_inference_steps=50,
+            audio_length_in_s=duration
+        ).audios[0]
+        
+        # Save audio to a file
+        output_path = "output.wav"
+        sf.write(output_path, audio, 22050)  # 22kHz sample rate
+        logger.info("Music generation complete.")
+        return output_path
+
+    # Gradio interface
+    logger.info("Setting up Gradio interface...")
+    with gr.Blocks(title="PhantomStep: Text-to-Music Generation 🎵") as demo:
+        gr.Markdown("# PhantomStep by GhostAI 🚀")
+        gr.Markdown("Enter a text prompt to generate music! 🎶")
+        
+        prompt_input = gr.Textbox(label="Text Prompt", placeholder="A jazzy piano melody with a fast tempo")
+        duration_input = gr.Slider(label="Duration (seconds)", minimum=10, maximum=60, value=20, step=1)
+        seed_input = gr.Number(label="Random Seed", value=42, precision=0)
+        generate_button = gr.Button("Generate Music")
+        
+        audio_output = gr.Audio(label="Generated Music")
+        
+        generate_button.click(
+            fn=generate_music,
+            inputs=[prompt_input, duration_input, seed_input],
+            outputs=audio_output
+        )
+
+    logger.info("Launching Gradio app...")
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
 
-demo.launch()
+except Exception as e:
+    logger.error(f"Failed to start the application: {str(e)}")
+    raise