Update app.py

app.py

@@ -4,7 +4,6 @@ from diffusers import DiffusionPipeline
 import torch
 import re
 import time
-import gc
 import soundfile as sf
 from qwen_tts import Qwen3TTSModel
 from langdetect import detect
@@ -12,19 +11,18 @@ import os
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- Chat ---
+# --- Chat Logic ---
 chat_models = {
     "Normal": "Qwen/Qwen3-0.6B",
     "Thinking": "Qwen/Qwen2.5-1.5B-Instruct"
 }
 loaded_chat_models = {}
 loaded_chat_tokenizers = {}
-chat_model_loaded = {}
 
 def load_chat_model(mode):
     model_id = chat_models[mode]
     if model_id not in loaded_chat_models:
-        gr.Info(f"🟡 Loading {mode} Brain...")
+        gr.Info(f"Loading {mode} Brain...")
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
@@ -33,204 +31,142 @@ def load_chat_model(mode):
         )
         loaded_chat_models[model_id] = model
         loaded_chat_tokenizers[model_id] = tokenizer
-        chat_model_loaded[model_id] = True
-        return "🟢 Model Loaded"
-    return "🟢 Model Already Loaded"
+        return "Status: 🟢 Loaded"
+    return "Status: 🟢 Loaded"
 
 def chat_logic(user_input, mode):
     model_id = chat_models[mode]
-    if model_id not in chat_model_loaded:
-        return "❌ Model Not Loaded"
+    if model_id not in loaded_chat_models:
+        return "❌ Please click 'Load Model' first!"
     model, tokenizer = loaded_chat_models[model_id], loaded_chat_tokenizers[model_id]
     messages = [{"role": "user", "content": user_input}]
     text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     generated_ids = model.generate(**model_inputs, max_new_tokens=1024)
-    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    cleaned_response = re.sub(r'<think>.*?</think>\s*\n?', '', response, flags=re.DOTALL)
-    return cleaned_response.strip()
-
-def clear_chat_model(password):
-    if password != "Creeper":
-        return "❌ Incorrect password"
-    for model_id in list(loaded_chat_models.keys()):
-        del loaded_chat_models[model_id]
-        del loaded_chat_tokenizers[model_id]
-        chat_model_loaded.pop(model_id, None)
-    
-    gc.collect()
-    torch.cuda.empty_cache()
-    if torch.cuda.is_available():
-        torch.cuda.ipc_collect()
-    return "🔴 Model Not Loaded (RAM Flushed)"
-
-# --- Image ---
+    response = tokenizer.batch_decode([generated_ids[0][len(model_inputs.input_ids[0]):]], skip_special_tokens=True)[0]
+    return re.sub(r'<think>.*?</think>\s*\n?', '', response, flags=re.DOTALL).strip()
+
+# --- Image Logic ---
 image_model_id = "stabilityai/sdxl-turbo"
 image_pipe = None
-image_model_loaded = False
 
 def load_image_model():
-    global image_pipe, image_model_loaded
+    global image_pipe
     if image_pipe is None:
-        gr.Info("🟡 Priming TNT (Image Gen)...")
+        gr.Info("Priming TNT Engine...")
         pipe = DiffusionPipeline.from_pretrained(
             image_model_id,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32
         )
         pipe.to(device)
         image_pipe = pipe
-        image_model_loaded = True
-        return "🟢 Model Loaded"
-    return "🟢 Model Already Loaded"
+        return "Status: 🟢 Loaded"
+    return "Status: 🟢 Loaded"
 
 def image_logic(prompt, width, height, steps):
-    if not image_model_loaded or image_pipe is None:
-        yield "❌ Model Not Loaded", None
+    if image_pipe is None:
+        yield "❌ Load Model First", None
         return
-    start_time = time.time()
-    final_prompt = f"{prompt}, centered and realistic (if applicable)"
     yield "💥 IGNITING...", None
     image = image_pipe(
-        prompt=final_prompt,
-        width=int(width),
-        height=int(height),
-        num_inference_steps=int(steps),
-        guidance_scale=0.0,
-        output_type="pil"
+        prompt=prompt, 
+        width=int(width), 
+        height=int(height), 
+        num_inference_steps=int(steps), 
+        guidance_scale=0.0
     ).images[0]
-    duration = round(time.time() - start_time, 2)
-    yield f"💥 EXPLODED in {duration}s", image
-
-def clear_image_model(password):
-    global image_pipe, image_model_loaded
-    if password != "Creeper":
-        return "❌ Incorrect password"
-    if image_pipe:
-        del image_pipe
-        image_pipe = None
-        image_model_loaded = False
-        gc.collect()
-        torch.cuda.empty_cache()
-        if torch.cuda.is_available():
-            torch.cuda.ipc_collect()
-    return "🔴 Model Not Loaded (RAM Flushed)"
-
-# --- TTS ---
+    yield "💥 EXPLODED", image
+
+# --- TTS Logic ---
 tts_model_id = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 SUPPORTED_VOICES = ['aiden', 'dylan', 'eric', 'ono_anna', 'ryan', 'serena', 'sohee', 'uncle_fu', 'vivian']
 tts_model = None
-tts_model_loaded = False
 
 def load_tts_model():
-    global tts_model, tts_model_loaded
+    global tts_model
     if tts_model is None:
-        gr.Info("🟡 Tuning Note-Blocks (TTS)...")
+        gr.Info("Tuning Note-Blocks...")
         tts_model = Qwen3TTSModel.from_pretrained(
             tts_model_id,
             device_map=device,
             torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
         )
-        tts_model_loaded = True
-        return "🟢 Model Loaded"
-    return "🟢 Model Already Loaded"
+        return "Status: 🟢 Loaded"
+    return "Status: 🟢 Loaded"
 
-def tts_logic(text, voice, instructions, auto_detect):
-    if not tts_model_loaded or tts_model is None:
-        return None, "❌ Model Not Loaded"
+def tts_logic(text, voice, inst, auto):
+    if tts_model is None:
+        return None, "Status: ❌ Not Loaded"
     try:
-        lang_map = {'zh': 'Chinese', 'en': 'English', 'jp': 'Japanese', 'ko': 'Korean'}
-        detected_lang = "English"
-        if auto_detect:
-            try:
-                raw_lang = detect(text).split('-')[0]
-                detected_lang = lang_map.get(raw_lang, "English")
-            except: pass
         wavs, sr = tts_model.generate_custom_voice(
-            language=detected_lang, speaker=voice, instruct=instructions, text=text
+            language="English", 
+            speaker=voice, 
+            instruct=inst, 
+            text=text
         )
         output_path = "creeper_voice.wav"
         sf.write(output_path, wavs[0], sr)
-        return output_path, f"Speaker: {voice} | Lang: {detected_lang}"
+        return output_path, "Status: 🟢 Audio Generated"
     except Exception as e:
-        return None, f"Error: {str(e)}"
-
-def clear_tts_model(password):
-    global tts_model, tts_model_loaded
-    if password != "Creeper":
-        return "❌ Incorrect password"
-    if tts_model:
-        del tts_model
-        tts_model = None
-        tts_model_loaded = False
-        gc.collect()
-        torch.cuda.empty_cache()
-        if torch.cuda.is_available():
-            torch.cuda.ipc_collect()
-    return "🔴 Model Not Loaded (RAM Flushed)"
-
-# --- UI ---
+        return None, f"Status: ❌ Error: {str(e)}"
+
+# --- UI Styles ---
 creeper_css = """
 body { background-color: #000000; }
-.gradio-container { background-color: #1e1e1e; border: 10px solid #2e8b57 !important; font-family: 'Courier New', Courier, monospace; color: #00ff00; }
-footer { display: none !important; }
-.gr-button-primary { background-color: #4A7023 !important; border: 4px solid #000 !important; color: white !important; font-weight: bold; text-transform: uppercase; }
-label span { color: #2e8b57 !important; font-weight: bold; font-size: 1.2em; }
-textarea, input, select, .gr-dropdown { background-color: #2e2e2e !important; color: #00ff00 !important; border: 3px solid #4A7023 !important; }
+.gradio-container { background-color: #1e1e1e; border: 8px solid #2e8b57; color: #00ff00; }
+.small-status { font-size: 0.8em; color: #2e8b57; margin-top: -10px; }
+.gr-button-primary { background-color: #4A7023 !important; border: 2px solid #000 !important; }
 """
 
 with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
     gr.Markdown("# 🟩 CREEPER AI HUB 🟩")
+
     with gr.Tabs():
+        # --- Chat Tab ---
         with gr.TabItem("SSSSS-CHAT"):
-            chat_status = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
-                mode_radio = gr.Radio(["Normal", "Thinking"], value="Normal", label="Mode")
-                load_chat_btn = gr.Button("Load Chat")
-                chat_pw = gr.Textbox(label="Password", type="password")
-                clear_chat_btn = gr.Button("Clear RAM")
-            chat_input = gr.Textbox(label="Message")
-            chat_output = gr.Textbox(label="Creeper Says")
+                chat_status = gr.Markdown("Status: 🔴 Not Loaded", elem_classes=["small-status"])
+            with gr.Row():
+                mode_radio = gr.Radio(["Normal", "Thinking"], value="Normal", label="Brain Mode")
+                load_chat_btn = gr.Button("Load Model")
+            chat_in = gr.Textbox(label="Message", lines=3)
+            chat_out = gr.Textbox(label="Creeper Says")
             chat_btn = gr.Button("EXPLODE TEXT", variant="primary")
+            
             load_chat_btn.click(load_chat_model, mode_radio, chat_status)
-            chat_btn.click(chat_logic, [chat_input, mode_radio], chat_output)
-            clear_chat_btn.click(clear_chat_model, chat_pw, chat_status)
+            chat_btn.click(chat_logic, [chat_in, mode_radio], chat_out)
 
+        # --- Image Tab ---
         with gr.TabItem("TNT-IMAGE"):
-            img_status = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
-                with gr.Column():
-                    img_prompt = gr.Textbox(label="Prompt")
-                    w_s = gr.Slider(256, 1024, 512, step=64, label="Width")
-                    h_s = gr.Slider(256, 1024, 512, step=64, label="Height")
-                    s_s = gr.Slider(1, 10, 4, step=1, label="Steps")
-                    load_img_btn = gr.Button("Load Image")
-                    img_btn = gr.Button("EXPLODE IMAGE", variant="primary")
-                    img_pw = gr.Textbox(label="Password", type="password")
-                    clear_img_btn = gr.Button("Clear RAM")
-                img_out = gr.Image(label="Loot")
+                img_status = gr.Markdown("Status: 🔴 Not Loaded", elem_classes=["small-status"])
+            with gr.Row():
+                img_prompt = gr.Textbox(label="Visual Idea", placeholder="Pixel art forest...")
+                load_img_btn = gr.Button("Load Model")
+            with gr.Row():
+                w_s = gr.Slider(256, 1024, 512, step=64, label="Width")
+                h_s = gr.Slider(256, 1024, 512, step=64, label="Height")
+                s_s = gr.Slider(1, 10, 4, step=1, label="Steps")
+            img_btn = gr.Button("EXPLODE IMAGE", variant="primary")
+            img_out = gr.Image(label="Output")
+            
             load_img_btn.click(load_image_model, None, img_status)
             img_btn.click(image_logic, [img_prompt, w_s, h_s, s_s], [img_status, img_out])
-            clear_img_btn.click(clear_image_model, img_pw, img_status)
 
-        with gr.TabItem("NOTE-BLOCK (TTS)"):
-            tts_status = gr.Label("🔴 Model Not Loaded", label="Status")
+        # --- TTS Tab ---
+        with gr.TabItem("NOTE-BLOCK"):
+            with gr.Row():
+                tts_status = gr.Markdown("Status: 🔴 Not Loaded", elem_classes=["small-status"])
             with gr.Row():
-                with gr.Column():
-                    tts_in = gr.Textbox(label="Text")
-                    voice_sel = gr.Dropdown(SUPPORTED_VOICES, value="vivian", label="Voice")
-                    auto_l = gr.Checkbox(True, label="Auto-detect")
-                    style_in = gr.Textbox("Speak naturally", label="Style")
-                    load_tts_btn = gr.Button("Load TTS")
-                    tts_btn = gr.Button("EXPLODE AUDIO", variant="primary")
-                    tts_pw = gr.Textbox(label="Password", type="password")
-                    clear_tts_btn = gr.Button("Clear RAM")
-                with gr.Column():
-                    aud_out = gr.Audio(label="Audio", type="filepath")
-                    meta_out = gr.Label(label="Metadata")
+                voice_sel = gr.Dropdown(SUPPORTED_VOICES, value="vivian", label="Speaker")
+                load_tts_btn = gr.Button("Load Model")
+            tts_in = gr.Textbox(label="Text to Speak")
+            style_in = gr.Textbox("Speak naturally", label="Instructions")
+            tts_btn = gr.Button("EXPLODE AUDIO", variant="primary")
+            aud_out = gr.Audio(label="Audio Output")
+            
             load_tts_btn.click(load_tts_model, None, tts_status)
-            tts_btn.click(tts_logic, [tts_in, voice_sel, style_in, auto_l], [aud_out, meta_out])
-            clear_tts_btn.click(clear_tts_model, tts_pw, tts_status)
+            tts_btn.click(tts_logic, [tts_in, voice_sel, style_in, gr.State(True)], [aud_out, tts_status])
 
 if __name__ == "__main__":
     demo.launch()