Update app.py

app.py

@@ -11,46 +11,44 @@ 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 = {}
-chat_status_label = None
 chat_model_loaded = {}
 
 def load_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...")
+        gr.Info(f"🟡 Loading {mode} Brain...")
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            torch_dtype="auto",
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
             device_map="auto"
         )
         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 "✅ Chat model loaded"
+        return "🟢 Model Loaded"
+    return "🟢 Model Already Loaded"
 
 def chat_logic(user_input, mode):
     model_id = chat_models[mode]
     if model_id not in chat_model_loaded:
-        return "❌ Model Not Loaded"
+        return "❌ Model Not Loaded. Click 'Load Chat 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()
 
@@ -62,21 +60,16 @@ def clear_chat_model(password):
         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")
-    return "✅ Chat model cleared from RAM"
+    return "🔴 Model Not Loaded"
 
-# --- Image ---
 image_model_id = "stabilityai/sdxl-turbo"
 image_pipe = None
-image_status_label = None
 image_model_loaded = False
 
 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...")
+        gr.Info("🟡 Priming TNT (Loading Image Model)...")
         pipe = DiffusionPipeline.from_pretrained(
             image_model_id,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32
@@ -84,25 +77,27 @@ def load_image_model():
         pipe.to(device)
         image_pipe = pipe
         image_model_loaded = True
-        if image_status_label:
-            image_status_label.update("🟢 Model Loaded")
-    return "✅ Image model loaded"
+        return "🟢 Model Loaded"
+    return "🟢 Model Already Loaded"
 
 def image_logic(prompt, width, height, steps):
     if not image_model_loaded or image_pipe is None:
-        return "❌ Model Not Loaded", 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
+    
     image = image_pipe(
         prompt=final_prompt,
         width=int(width),
         height=int(height),
         num_inference_steps=int(steps),
-        guidance_scale=8.0,
-        lcm_origin_steps=50,
+        guidance_scale=0.0,
         output_type="pil"
     ).images[0]
+    
     duration = round(time.time() - start_time, 2)
     yield f"💥 EXPLODED in {duration}s", image
 
@@ -115,31 +110,25 @@ def clear_image_model(password):
         image_pipe = None
         image_model_loaded = False
         torch.cuda.empty_cache()
-    if image_status_label:
-        image_status_label.update("🔴 Model Not Loaded")
-    return "✅ Image model cleared from RAM"
+    return "🔴 Model Not Loaded"
 
-# --- 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
 tts_model_loaded = False
 
 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...")
+        gr.Info("🟡 Tuning Note-Blocks (Loading TTS)...")
         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
-        if tts_status_label:
-            tts_status_label.update("🟢 Model Loaded")
-    return "✅ TTS model loaded"
+        return "🟢 Model Loaded"
+    return "🟢 Model Already Loaded"
 
 def tts_logic(text, voice, instructions, auto_detect):
     if not tts_model_loaded or tts_model is None:
@@ -157,6 +146,7 @@ def tts_logic(text, voice, instructions, auto_detect):
                 detected_lang = lang_map.get(raw_lang, "English")
             except:
                 pass
+        
         wavs, sr = tts_model.generate_custom_voice(
             language=detected_lang,
             speaker=voice,
@@ -178,11 +168,8 @@ def clear_tts_model(password):
         tts_model = None
         tts_model_loaded = False
         torch.cuda.empty_cache()
-    if tts_status_label:
-        tts_status_label.update("🔴 Model Not Loaded")
-    return "✅ TTS model cleared from RAM"
+    return "🔴 Model Not Loaded"
 
-# --- 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; }
@@ -199,23 +186,23 @@ 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")
                 load_chat_btn = gr.Button("Load Chat Model")
-                clear_chat_btn = gr.Button("Clear Model (Password Protected)")
+                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=gr.Textbox(label="Enter Password"), outputs=chat_status_label)
+            clear_chat_btn.click(fn=clear_chat_model, inputs=chat_pw, 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")
@@ -223,38 +210,41 @@ with gr.Blocks(css=creeper_css, title="CREEPER AI HUB") as demo:
                 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, 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")
+                        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")
                     img_btn = gr.Button("EXPLODE IMAGE", variant="primary")
-                    clear_image_btn = gr.Button("Clear Model (Password Protected)")
+                    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=gr.Textbox(label="Enter Password"), outputs=image_status_label)
+            clear_image_btn.click(fn=clear_image_model, inputs=img_pw, 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)
+                    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_btn = gr.Button("EXPLODE AUDIO", variant="primary")
-                    clear_tts_btn = gr.Button("Clear Model (Password Protected)")
+                    tts_pw = gr.Textbox(label="Password", type="password")
+                    clear_tts_btn = gr.Button("Clear Model")
                 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=gr.Textbox(label="Enter Password"), outputs=tts_status_label)
+            clear_tts_btn.click(fn=clear_tts_model, inputs=tts_pw, outputs=tts_status_label)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()