Update app.py

app.py

@@ -19,8 +19,9 @@ chat_models = {
 loaded_chat_models = {}
 loaded_chat_tokenizers = {}
 chat_status_label = None
+chat_model_loaded = {}
 
-def get_chat_model(mode):
+def load_chat_model(mode):
     global chat_status_label
     model_id = chat_models[mode]
     if model_id not in loaded_chat_models:
@@ -34,12 +35,16 @@ def get_chat_model(mode):
         )
         loaded_chat_models[model_id] = model
         loaded_chat_tokenizers[model_id] = tokenizer
+        chat_model_loaded[model_id] = True
         if chat_status_label:
             chat_status_label.update("🟢 Model Loaded")
-    return loaded_chat_models[model_id], loaded_chat_tokenizers[model_id]
+    return "✅ Chat model loaded"
 
 def chat_logic(user_input, mode):
-    model, tokenizer = get_chat_model(mode)
+    model_id = chat_models[mode]
+    if model_id not in chat_model_loaded:
+        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)
@@ -55,6 +60,7 @@ def clear_chat_model(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)
     torch.cuda.empty_cache()
     if chat_status_label:
         chat_status_label.update("🔴 Model Not Loaded")
@@ -64,9 +70,10 @@ def clear_chat_model(password):
 image_model_id = "stabilityai/sdxl-turbo"
 image_pipe = None
 image_status_label = None
+image_model_loaded = False
 
-def get_image_model():
-    global image_pipe
+def load_image_model():
+    global image_pipe, image_model_loaded
     if image_pipe is None:
         if image_status_label:
             image_status_label.update("🟡 Model Loading...")
@@ -76,16 +83,18 @@ def get_image_model():
         )
         pipe.to(device)
         image_pipe = pipe
+        image_model_loaded = True
         if image_status_label:
             image_status_label.update("🟢 Model Loaded")
-    return image_pipe
+    return "✅ Image model loaded"
 
 def image_logic(prompt, width, height, steps):
-    pipe = get_image_model()
+    if not image_model_loaded or image_pipe is None:
+        return "❌ Model Not Loaded", None
     start_time = time.time()
     final_prompt = f"{prompt}, centered and realistic (if applicable)"
     yield "💥 IGNITING... (Image generator AI)...", None
-    image = pipe(
+    image = image_pipe(
         prompt=final_prompt,
         width=int(width),
         height=int(height),
@@ -98,12 +107,13 @@ def image_logic(prompt, width, height, steps):
     yield f"💥 EXPLODED in {duration}s", image
 
 def clear_image_model(password):
-    global image_pipe
+    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
         torch.cuda.empty_cache()
     if image_status_label:
         image_status_label.update("🔴 Model Not Loaded")
@@ -114,9 +124,10 @@ 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
+tts_model_loaded = False
 
-def get_tts_model():
-    global tts_model
+def load_tts_model():
+    global tts_model, tts_model_loaded
     if tts_model is None:
         if tts_status_label:
             tts_status_label.update("🟡 Model Loading...")
@@ -125,12 +136,14 @@ def get_tts_model():
             device_map=device,
             torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
         )
+        tts_model_loaded = True
         if tts_status_label:
             tts_status_label.update("🟢 Model Loaded")
-    return tts_model
+    return "✅ TTS model loaded"
 
 def tts_logic(text, voice, instructions, auto_detect):
-    model = get_tts_model()
+    if not tts_model_loaded or tts_model is None:
+        return None, "❌ Model Not Loaded"
     try:
         lang_map = {
             'zh': 'Chinese', 'en': 'English', 'jp': 'Japanese',
@@ -144,7 +157,7 @@ def tts_logic(text, voice, instructions, auto_detect):
                 detected_lang = lang_map.get(raw_lang, "English")
             except:
                 pass
-        wavs, sr = model.generate_custom_voice(
+        wavs, sr = tts_model.generate_custom_voice(
             language=detected_lang,
             speaker=voice,
             instruct=instructions,
@@ -157,12 +170,13 @@ def tts_logic(text, voice, instructions, auto_detect):
         return None, f"System Error: {str(e)}"
 
 def clear_tts_model(password):
-    global tts_model
+    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
         torch.cuda.empty_cache()
     if tts_status_label:
         tts_status_label.update("🔴 Model Not Loaded")
@@ -191,11 +205,13 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
             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")
+                load_chat_btn = gr.Button("Load Chat Model")
                 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")
+            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=gr.Textbox(label="Enter Password"), outputs=chat_status_label)
 
@@ -210,10 +226,12 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
                         w_slider = gr.Slider(256, 768, 512, step=64, label="Block Width")
                         h_slider = gr.Slider(256, 768, 512, step=64, label="Block Height")
                     s_slider = gr.Slider(4, 12, 5, step=1, label="Detonation Steps")
+                    load_image_btn = gr.Button("Load Image Model")
                     img_btn = gr.Button("EXPLODE IMAGE", variant="primary")
                     clear_image_btn = gr.Button("Clear Model (Password Protected)")
                 with gr.Column(scale=1):
                     img_output = gr.Image(label="Rendered Loot")
+            load_image_btn.click(fn=load_image_model, inputs=None, 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=gr.Textbox(label="Enter Password"), outputs=image_status_label)
 
@@ -228,11 +246,13 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
                         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_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")
+            load_tts_btn.click(fn=load_tts_model, inputs=None, 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=gr.Textbox(label="Enter Password"), outputs=tts_status_label)