Update app.py

app.py

@@ -11,29 +11,21 @@ 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"
 }
-
 loaded_chat_models = {}
 loaded_chat_tokenizers = {}
-
-image_model_id = "stabilityai/sdxl-turbo"
-image_pipe = DiffusionPipeline.from_pretrained(image_model_id)
-image_pipe.to(device)
-
-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 = Qwen3TTSModel.from_pretrained(
-    tts_model_id,
-    device_map=device,
-    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
-)
+chat_status_label = None
 
 def get_chat_model(mode):
+    global chat_status_label
     model_id = chat_models[mode]
     if model_id not in loaded_chat_models:
+        if chat_status_label:
+            chat_status_label.update("🟡 Model Loading...")
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
@@ -42,6 +34,8 @@ def get_chat_model(mode):
         )
         loaded_chat_models[model_id] = model
         loaded_chat_tokenizers[model_id] = tokenizer
+        if chat_status_label:
+            chat_status_label.update("🟢 Model Loaded")
     return loaded_chat_models[model_id], loaded_chat_tokenizers[model_id]
 
 def chat_logic(user_input, mode):
@@ -55,11 +49,43 @@ def chat_logic(user_input, mode):
     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]
+    torch.cuda.empty_cache()
+    if chat_status_label:
+        chat_status_label.update("🔴 Model Not Loaded")
+    return "✅ Chat model cleared from RAM"
+
+# --- Image ---
+image_model_id = "stabilityai/sdxl-turbo"
+image_pipe = None
+image_status_label = None
+
+def get_image_model():
+    global image_pipe
+    if image_pipe is None:
+        if image_status_label:
+            image_status_label.update("🟡 Model Loading...")
+        pipe = DiffusionPipeline.from_pretrained(
+            image_model_id,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        )
+        pipe.to(device)
+        image_pipe = pipe
+        if image_status_label:
+            image_status_label.update("🟢 Model Loaded")
+    return image_pipe
+
 def image_logic(prompt, width, height, steps):
+    pipe = get_image_model()
     start_time = time.time()
     final_prompt = f"{prompt}, centered and realistic (if applicable)"
     yield "💥 IGNITING... (Image generator AI)...", None
-    image = image_pipe(
+    image = pipe(
         prompt=final_prompt,
         width=int(width),
         height=int(height),
@@ -71,21 +97,54 @@ def image_logic(prompt, width, height, steps):
     duration = round(time.time() - start_time, 2)
     yield f"💥 EXPLODED in {duration}s", image
 
+def clear_image_model(password):
+    global image_pipe
+    if password != "Creeper":
+        return "❌ Incorrect password"
+    if image_pipe:
+        del image_pipe
+        image_pipe = None
+        torch.cuda.empty_cache()
+    if image_status_label:
+        image_status_label.update("🔴 Model Not Loaded")
+    return "✅ Image model cleared from RAM"
+
+# --- 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
+tts_status_label = None
+
+def get_tts_model():
+    global tts_model
+    if tts_model is None:
+        if tts_status_label:
+            tts_status_label.update("🟡 Model Loading...")
+        tts_model = Qwen3TTSModel.from_pretrained(
+            tts_model_id,
+            device_map=device,
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
+        )
+        if tts_status_label:
+            tts_status_label.update("🟢 Model Loaded")
+    return tts_model
+
 def tts_logic(text, voice, instructions, auto_detect):
+    model = get_tts_model()
     try:
         lang_map = {
-            'zh': 'Chinese', 'en': 'English', 'jp': 'Japanese', 
-            'ko': 'Korean', 'de': 'German', 'fr': 'French', 
+            'zh': 'Chinese', 'en': 'English', 'jp': 'Japanese',
+            'ko': 'Korean', 'de': 'German', 'fr': 'French',
             'ru': 'Russian', 'pt': 'Portuguese', 'es': 'Spanish', 'it': 'Italian'
         }
-        detected_lang = "English" 
+        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(
+        wavs, sr = model.generate_custom_voice(
             language=detected_lang,
             speaker=voice,
             instruct=instructions,
@@ -97,6 +156,19 @@ def tts_logic(text, voice, instructions, auto_detect):
     except Exception as e:
         return None, f"System Error: {str(e)}"
 
+def clear_tts_model(password):
+    global tts_model
+    if password != "Creeper":
+        return "❌ Incorrect password"
+    if tts_model:
+        del tts_model
+        tts_model = None
+        torch.cuda.empty_cache()
+    if tts_status_label:
+        tts_status_label.update("🔴 Model Not Loaded")
+    return "✅ TTS model cleared from RAM"
+
+# --- 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; }
@@ -111,20 +183,26 @@ textarea, input, .gr-box, .gr-input, select, .gr-dropdown { background-color: #2
 
 with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
     gr.Markdown("# 🟩 CREEPER AI HUB 🟩")
-    
+
     with gr.Tabs():
+        # --- Chat ---
         with gr.TabItem("SSSSS-CHAT"):
             gr.Markdown("### Qwen Chat System")
+            chat_status_label = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
                 mode_radio = gr.Radio(choices=["Normal", "Thinking"], value="Normal", label="Select Brain Mode")
+                clear_chat_btn = gr.Button("Clear Model (Password Protected)")
             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")
             chat_btn.click(fn=chat_logic, inputs=[chat_input, mode_radio], outputs=chat_output)
-            
+            clear_chat_btn.click(fn=clear_chat_model, inputs=gr.Textbox(label="Enter Password"), outputs=chat_status_label)
+
+        # --- Image ---
         with gr.TabItem("TNT-IMAGE"):
             gr.Markdown("### Image Generator System")
+            image_status_label = 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)
@@ -133,13 +211,16 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
                         h_slider = gr.Slider(256, 768, 512, step=64, label="Block Height")
                     s_slider = gr.Slider(4, 12, 5, step=1, label="Detonation Steps")
                     img_btn = gr.Button("EXPLODE IMAGE", variant="primary")
+                    clear_image_btn = gr.Button("Clear Model (Password Protected)")
                 with gr.Column(scale=1):
-                    img_status = gr.Markdown("### Status: 🟢 Armed")
                     img_output = gr.Image(label="Rendered Loot")
-            img_btn.click(fn=image_logic, inputs=[img_prompt, w_slider, h_slider, s_slider], outputs=[img_status, img_output])
+            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=gr.Textbox(label="Enter Password"), outputs=image_status_label)
 
+        # --- TTS ---
         with gr.TabItem("NOTE-BLOCK (TTS)"):
             gr.Markdown("### Smart Audio Studio")
+            tts_status_label = gr.Label("🔴 Model Not Loaded", label="Status")
             with gr.Row():
                 with gr.Column():
                     tts_input = gr.Textbox(label="Text to Speak", placeholder="Enter text for the block to play...", lines=4)
@@ -148,10 +229,12 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
                         auto_lang = gr.Checkbox(label="Auto-detect Language", value=True)
                     style_instruct = gr.Textbox(label="Style Instruction", value="Speak naturally")
                     tts_btn = gr.Button("EXPLODE AUDIO", variant="primary")
+                    clear_tts_btn = gr.Button("Clear Model (Password Protected)")
                 with gr.Column():
                     audio_output = gr.Audio(label="Audio Output", type="filepath")
                     status_info = gr.Label(label="Block Metadata")
             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=gr.Textbox(label="Enter Password"), outputs=tts_status_label)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()