Update app.py

app.py

@@ -4,6 +4,7 @@ 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
@@ -11,6 +12,7 @@ import os
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# --- Chat ---
 chat_models = {
     "Normal": "Qwen/Qwen3-0.6B",
     "Thinking": "Qwen/Qwen2.5-1.5B-Instruct"
@@ -38,17 +40,14 @@ def load_chat_model(mode):
 def chat_logic(user_input, mode):
     model_id = chat_models[mode]
     if model_id not in chat_model_loaded:
-        return "❌ Model Not Loaded. Click 'Load Chat Model' first!"
-    
+        return "❌ Model Not Loaded"
     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()
 
@@ -59,9 +58,14 @@ def clear_chat_model(password):
         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()
-    return "🔴 Model Not Loaded"
+    if torch.cuda.is_available():
+        torch.cuda.ipc_collect()
+    return "🔴 Model Not Loaded (RAM Flushed)"
 
+# --- Image ---
 image_model_id = "stabilityai/sdxl-turbo"
 image_pipe = None
 image_model_loaded = False
@@ -69,7 +73,7 @@ image_model_loaded = False
 def load_image_model():
     global image_pipe, image_model_loaded
     if image_pipe is None:
-        gr.Info("🟡 Priming TNT (Loading Image Model)...")
+        gr.Info("🟡 Priming TNT (Image Gen)...")
         pipe = DiffusionPipeline.from_pretrained(
             image_model_id,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32
@@ -84,11 +88,9 @@ def image_logic(prompt, width, height, steps):
     if not image_model_loaded or image_pipe is None:
         yield "❌ Model Not Loaded", None
         return
-
     start_time = time.time()
     final_prompt = f"{prompt}, centered and realistic (if applicable)"
-    yield "💥 IGNITING... (Image generator AI)...", None
-    
+    yield "💥 IGNITING...", None
     image = image_pipe(
         prompt=final_prompt,
         width=int(width),
@@ -97,7 +99,6 @@ def image_logic(prompt, width, height, steps):
         guidance_scale=0.0,
         output_type="pil"
     ).images[0]
-    
     duration = round(time.time() - start_time, 2)
     yield f"💥 EXPLODED in {duration}s", image
 
@@ -109,9 +110,13 @@ def clear_image_model(password):
         del image_pipe
         image_pipe = None
         image_model_loaded = False
+        gc.collect()
         torch.cuda.empty_cache()
-    return "🔴 Model Not Loaded"
+        if torch.cuda.is_available():
+            torch.cuda.ipc_collect()
+    return "🔴 Model Not Loaded (RAM Flushed)"
 
+# --- TTS ---
 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
@@ -120,7 +125,7 @@ tts_model_loaded = False
 def load_tts_model():
     global tts_model, tts_model_loaded
     if tts_model is None:
-        gr.Info("🟡 Tuning Note-Blocks (Loading TTS)...")
+        gr.Info("🟡 Tuning Note-Blocks (TTS)...")
         tts_model = Qwen3TTSModel.from_pretrained(
             tts_model_id,
             device_map=device,
@@ -134,30 +139,21 @@ def tts_logic(text, voice, instructions, auto_detect):
     if not tts_model_loaded or tts_model is None:
         return None, "❌ Model Not Loaded"
     try:
-        lang_map = {
-            'zh': 'Chinese', 'en': 'English', 'jp': 'Japanese',
-            'ko': 'Korean', 'de': 'German', 'fr': 'French',
-            'ru': 'Russian', 'pt': 'Portuguese', 'es': 'Spanish', 'it': 'Italian'
-        }
+        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
-        
+            except: pass
         wavs, sr = tts_model.generate_custom_voice(
-            language=detected_lang,
-            speaker=voice,
-            instruct=instructions,
-            text=text
+            language=detected_lang, speaker=voice, instruct=instructions, text=text
         )
         output_path = "creeper_voice.wav"
         sf.write(output_path, wavs[0], sr)
-        return output_path, f"Language: {detected_lang} | Speaker: {voice}"
+        return output_path, f"Speaker: {voice} | Lang: {detected_lang}"
     except Exception as e:
-        return None, f"System Error: {str(e)}"
+        return None, f"Error: {str(e)}"
 
 def clear_tts_model(password):
     global tts_model, tts_model_loaded
@@ -167,84 +163,74 @@ def clear_tts_model(password):
         del tts_model
         tts_model = None
         tts_model_loaded = False
+        gc.collect()
         torch.cuda.empty_cache()
-    return "🔴 Model Not Loaded"
+        if torch.cuda.is_available():
+            torch.cuda.ipc_collect()
+    return "🔴 Model Not Loaded (RAM Flushed)"
 
+# --- UI ---
 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; }
-.gr-button-primary:hover { background-color: #5ea032 !important; box-shadow: 0 0 20px #2e8b57; }
 label span { color: #2e8b57 !important; font-weight: bold; font-size: 1.2em; }
-textarea, input, .gr-box, .gr-input, select, .gr-dropdown { background-color: #2e2e2e !important; color: #00ff00 !important; border: 3px solid #4A7023 !important; }
-.tabs { border-bottom: 5px solid #4A7023 !important; }
-.tab-nav button.selected { background-color: #4A7023 !important; color: white !important; }
+textarea, input, select, .gr-dropdown { background-color: #2e2e2e !important; color: #00ff00 !important; border: 3px solid #4A7023 !important; }
 """
 
 with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
     gr.Markdown("# 🟩 CREEPER AI HUB 🟩")
-
     with gr.Tabs():
         with gr.TabItem("SSSSS-CHAT"):
-            gr.Markdown("### Qwen Chat System")
-            chat_status_label = gr.Label("🔴 Model Not Loaded", label="Status")
+            chat_status = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
-                mode_radio = gr.Radio(choices=["Normal", "Thinking"], value="Normal", label="Select Brain Mode")
-                load_chat_btn = gr.Button("Load Chat Model")
+                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 Model")
-            with gr.Column():
-                chat_input = gr.Textbox(lines=4, placeholder="Ssssss... Talk to the Creeper...", label="Message")
-                chat_output = gr.Textbox(label="Creeper Says")
-                chat_btn = gr.Button("EXPLODE TEXT", variant="primary")
-            
-            load_chat_btn.click(fn=load_chat_model, inputs=mode_radio, outputs=chat_status_label)
-            chat_btn.click(fn=chat_logic, inputs=[chat_input, mode_radio], outputs=chat_output)
-            clear_chat_btn.click(fn=clear_chat_model, inputs=chat_pw, outputs=chat_status_label)
+                clear_chat_btn = gr.Button("Clear RAM")
+            chat_input = gr.Textbox(label="Message")
+            chat_output = 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)
 
         with gr.TabItem("TNT-IMAGE"):
-            gr.Markdown("### Image Generator System")
-            image_status_label = gr.Label("🔴 Model Not Loaded", label="Status")
+            img_status = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
-                with gr.Column(scale=1):
-                    img_prompt = gr.Textbox(label="Visual Idea", placeholder="A pixelated forest...", lines=3)
-                    with gr.Row():
-                        w_slider = gr.Slider(256, 1024, 512, step=64, label="Block Width")
-                        h_slider = gr.Slider(256, 1024, 512, step=64, label="Block Height")
-                    s_slider = gr.Slider(1, 10, 4, step=1, label="Detonation Steps")
-                    load_image_btn = gr.Button("Load Image Model")
+                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_image_btn = gr.Button("Clear Model")
-                with gr.Column(scale=1):
-                    img_output = gr.Image(label="Rendered Loot")
-            
-            load_image_btn.click(fn=load_image_model, inputs=[], outputs=image_status_label)
-            img_btn.click(fn=image_logic, inputs=[img_prompt, w_slider, h_slider, s_slider], outputs=[image_status_label, img_output])
-            clear_image_btn.click(fn=clear_image_model, inputs=img_pw, outputs=image_status_label)
+                    clear_img_btn = gr.Button("Clear RAM")
+                img_out = gr.Image(label="Loot")
+            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)"):
-            gr.Markdown("### Smart Audio Studio")
-            tts_status_label = gr.Label("🔴 Model Not Loaded", label="Status")
+            tts_status = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
                 with gr.Column():
-                    tts_input = gr.Textbox(label="Text to Speak", placeholder="Enter text...", lines=4)
-                    with gr.Row():
-                        voice_select = gr.Dropdown(choices=SUPPORTED_VOICES, value="vivian", label="Select Speaker")
-                        auto_lang = gr.Checkbox(label="Auto-detect Language", value=True)
-                    style_instruct = gr.Textbox(label="Style Instruction", value="Speak naturally")
-                    load_tts_btn = gr.Button("Load TTS Model")
+                    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 Model")
+                    clear_tts_btn = gr.Button("Clear RAM")
                 with gr.Column():
-                    audio_output = gr.Audio(label="Audio Output", type="filepath")
-                    status_info = gr.Label(label="Block Metadata")
-            
-            load_tts_btn.click(fn=load_tts_model, inputs=[], outputs=tts_status_label)
-            tts_btn.click(fn=tts_logic, inputs=[tts_input, voice_select, style_instruct, auto_lang], outputs=[audio_output, status_info])
-            clear_tts_btn.click(fn=clear_tts_model, inputs=tts_pw, outputs=tts_status_label)
+                    aud_out = gr.Audio(label="Audio", type="filepath")
+                    meta_out = gr.Label(label="Metadata")
+            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)
 
 if __name__ == "__main__":
     demo.launch()