Update app.py

app.py

@@ -4,186 +4,359 @@ import imageio
 import torch
 import time
 import os
-from openai import OpenAI
-from transformers import pipeline
+
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 from diffusers import DiffusionPipeline
 
-# ---------- Load OpenAI Key from Hugging Face Secret ----------
-client = OpenAI(api_key=os.getenv("OPENAI_KEY"))
 
-# ---------- Configuration ----------
-AVAILABLE_MODELS = {
-    "Codette Fine-Tuned (v9)": "ft:gpt-4.1-2025-04-14:raiffs-bits:codette-final:BO907H7Z",
-    "GPT-2 (small, fast)": "gpt2",
-    "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
-    "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
-}
+# -------------------------------------------------
+# Device Setup
+# -------------------------------------------------
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
 text_model_cache = {}
 chat_memory = {}
-last_usage_time = {}
 
-MAX_PROMPTS_PER_SESSION = 5
-THROTTLE_SECONDS = 30
 
-# ---------- Load Image Generator ----------
+# -------------------------------------------------
+# Available Text Models
+# -------------------------------------------------
+
+AVAILABLE_MODELS = {
+    "Codette LoRA (Llama-3.1)": "codette_lora",
+    "Mistral-7B Instruct": "mistralai/Mistral-7B-Instruct-v0.2",
+    "Phi-3 Mini": "microsoft/phi-3-mini-4k-instruct",
+    "GPT-2 (lightweight)": "gpt2"
+}
+
+
+# -------------------------------------------------
+# Load Codette LoRA Adapter
+# -------------------------------------------------
+
+def load_codette_lora():
+
+    base_model = "meta-llama/Meta-Llama-3.1-8B"
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        device_map="auto"
+    )
+
+    model = PeftModel.from_pretrained(
+        model,
+        "Raiff1982/codette-lora-adapters"
+    )
+
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device_map="auto"
+    )
+
+    return pipe
+
+
+# -------------------------------------------------
+# Image Generator (SDXL Turbo)
+# -------------------------------------------------
+
 try:
+
     image_generator = DiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        safety_checker=None,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        variant="fp16" if device == "cuda" else None
     )
+
     image_generator.to(device)
+
     image_enabled = True
+
 except Exception as e:
+
     print(f"[Image Model Load Error]: {e}")
     image_generator = None
     image_enabled = False
 
-# ---------- Load Video Generator ----------
+
+# -------------------------------------------------
+# Video Generator (Zeroscope)
+# -------------------------------------------------
+
 try:
+
     video_pipeline = DiffusionPipeline.from_pretrained(
-        "damo-vilab/text-to-video-ms-1.7b",
-        safety_checker=None,
+        "cerspense/zeroscope_v2_576w",
         torch_dtype=torch.float16 if device == "cuda" else torch.float32
     )
+
     video_pipeline.to(device)
+
     video_enabled = True
+
 except Exception as e:
+
     print(f"[Video Model Load Error]: {e}")
     video_pipeline = None
     video_enabled = False
 
-# ---------- Main Terminal with Rate Limits ----------
-def codette_terminal_limited(prompt, model_name, generate_image, generate_video,
-                              session_id, batch_size, video_steps, fps):
+
+# -------------------------------------------------
+# Load Text Models
+# -------------------------------------------------
+
+def get_text_model(model_name):
+
+    if model_name not in text_model_cache:
+
+        if AVAILABLE_MODELS[model_name] == "codette_lora":
+
+            text_model_cache[model_name] = load_codette_lora()
+
+        else:
+
+            text_model_cache[model_name] = pipeline(
+                "text-generation",
+                model=AVAILABLE_MODELS[model_name],
+                device=0 if device == "cuda" else -1
+            )
+
+    return text_model_cache[model_name]
+
+
+# -------------------------------------------------
+# Codette Terminal Logic
+# -------------------------------------------------
+
+def codette_terminal(prompt, model_name, generate_image, generate_video,
+                     session_id, batch_size, video_steps, fps):
+
     if session_id not in chat_memory:
         chat_memory[session_id] = []
 
     if prompt.lower() in ["exit", "quit"]:
         chat_memory[session_id] = []
-        yield "🧠 Codette signing off... Session reset.", None, None
+        yield "🧠 Codette session reset.", None, None
+        return
+
+    try:
+
+        model = get_text_model(model_name)
+
+        result = model(
+            prompt,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
+        )
+
+        output = result[0]["generated_text"]
+
+    except Exception as e:
+
+        yield f"[Text generation error]: {e}", None, None
         return
 
-    if model_name == "Codette Fine-Tuned (v9)":
-        count = sum(1 for line in chat_memory[session_id] if line.startswith("🖋️ You >"))
-        if count >= MAX_PROMPTS_PER_SESSION:
-            yield "[🛑 Limit] Max 5 prompts per session.", None, None
-            return
-        now = time.time()
-        if now - last_usage_time.get(session_id, 0) < THROTTLE_SECONDS:
-            wait = int(THROTTLE_SECONDS - (now - last_usage_time[session_id]))
-            yield f"[⏳ Wait] Try again in {wait} sec.", None, None
-            return
-        last_usage_time[session_id] = now
-
-    if model_name == "Codette Fine-Tuned (v9)":
-        try:
-            response = client.chat.completions.create(
-                model=AVAILABLE_MODELS[model_name],
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.7,
-                max_tokens=256
-            )
-            output = response.choices[0].message.content.strip()
-        except Exception as e:
-            yield f"[OpenAI error]: {e}", None, None
-            return
-    else:
-        if model_name not in text_model_cache:
-            try:
-                text_model_cache[model_name] = pipeline(
-                    "text-generation",
-                    model=AVAILABLE_MODELS[model_name],
-                    device=0 if device == "cuda" else -1
-                )
-            except Exception as e:
-                yield f"[Text model error]: {e}", None, None
-                return
-        try:
-            output = text_model_cache[model_name](
-                prompt, max_length=100, do_sample=True, num_return_sequences=1
-            )[0]['generated_text'].strip()
-        except Exception as e:
-            yield f"[Generation error]: {e}", None, None
-            return
 
-    # Stream text output
+    # -------------------------------------------------
+    # Stream Text
+    # -------------------------------------------------
+
     response_so_far = ""
+
     for char in output:
+
         response_so_far += char
+
         temp_log = chat_memory[session_id][:]
+
         temp_log.append(f"🖋️ You > {prompt}")
         temp_log.append(f"🧠 Codette > {response_so_far}")
+
         yield "\n".join(temp_log[-10:]), None, None
+
         time.sleep(0.01)
 
+
     chat_memory[session_id].append(f"🖋️ You > {prompt}")
     chat_memory[session_id].append(f"🧠 Codette > {output}")
 
-    imgs, vid = None, None
+
+    imgs = None
+    vid = None
+
+
+    # -------------------------------------------------
+    # Image Generation
+    # -------------------------------------------------
 
     if generate_image and image_enabled:
+
         try:
-            result = image_generator(prompt, num_images_per_prompt=batch_size)
+
+            result = image_generator(
+                prompt,
+                num_images_per_prompt=batch_size,
+                num_inference_steps=2
+            )
+
             imgs = result.images
+
         except Exception as e:
-            response_so_far += f"\n[Image error]: {e}"
+
+            print(f"[Image error]: {e}")
+
+
+    # -------------------------------------------------
+    # Video Generation
+    # -------------------------------------------------
 
     if generate_video and video_enabled:
+
         try:
-            result = video_pipeline(prompt, num_inference_steps=video_steps)
+
+            result = video_pipeline(
+                prompt,
+                num_inference_steps=video_steps
+            )
+
             frames = result.frames
-            temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+
+            temp_video_path = tempfile.NamedTemporaryFile(
+                suffix=".mp4",
+                delete=False
+            ).name
+
             imageio.mimsave(temp_video_path, frames, fps=fps)
+
             vid = temp_video_path
+
         except Exception as e:
-            response_so_far += f"\n[Video error]: {e}"
+
+            print(f"[Video error]: {e}")
+
 
     yield "\n".join(chat_memory[session_id][-10:]), imgs, vid
 
-# ---------- Gradio UI ----------
-with gr.Blocks(title="🧬 Codette Terminal – Streamed AI Chat") as demo:
-    gr.Markdown("## 🧬 Codette Terminal (Chat + Image + Video + Fine-Tuned AI)")
-    gr.Markdown("Type a prompt, choose a model, and generate responses. Type `'exit'` to reset the session.")
 
-    with gr.Row():
-        session_id = gr.Textbox(value="session_default", visible=False)
-        model_dropdown = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Language Model")
+# -------------------------------------------------
+# Gradio Interface
+# -------------------------------------------------
 
-    with gr.Row():
-        generate_image_toggle = gr.Checkbox(label="Generate Image(s)?", value=False, interactive=image_enabled)
-        generate_video_toggle = gr.Checkbox(label="Generate Video?", value=False, interactive=video_enabled)
+with gr.Blocks(title="🧬 Codette Terminal") as demo:
+
+    gr.Markdown("## 🧬 Codette Terminal")
+    gr.Markdown("Chat with Codette, generate images, and create short videos.")
+
+    session_id = gr.Textbox(value="default_session", visible=False)
 
     with gr.Row():
-        batch_size_slider = gr.Slider(label="Number of Images", minimum=1, maximum=4, step=1, value=1)
-        video_steps_slider = gr.Slider(label="Video Inference Steps", minimum=10, maximum=100, step=10, value=50)
-        fps_slider = gr.Slider(label="Video FPS", minimum=4, maximum=24, step=2, value=8)
+
+        model_dropdown = gr.Dropdown(
+            choices=list(AVAILABLE_MODELS.keys()),
+            value="Codette LoRA (Llama-3.1)",
+            label="Language Model"
+        )
 
     with gr.Row():
-        user_input = gr.Textbox(
-            label="Your Prompt",
-            placeholder="e.g. A robot dreaming on Mars",
-            lines=1
+
+        generate_image_toggle = gr.Checkbox(
+            label="Generate Image(s)",
+            value=False,
+            interactive=image_enabled
+        )
+
+        generate_video_toggle = gr.Checkbox(
+            label="Generate Video",
+            value=False,
+            interactive=video_enabled
         )
 
     with gr.Row():
-        output_text = gr.Textbox(label="Codette Output", lines=15, interactive=False)
+
+        batch_size_slider = gr.Slider(
+            label="Number of Images",
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=1
+        )
+
+        video_steps_slider = gr.Slider(
+            label="Video Inference Steps",
+            minimum=10,
+            maximum=50,
+            step=10,
+            value=20
+        )
+
+        fps_slider = gr.Slider(
+            label="Video FPS",
+            minimum=4,
+            maximum=24,
+            step=2,
+            value=8
+        )
+
+    user_input = gr.Textbox(
+        label="Your Prompt",
+        placeholder="A robot dreaming on Mars...",
+        lines=1
+    )
+
+    output_text = gr.Textbox(
+        label="Codette Output",
+        lines=15,
+        interactive=False
+    )
 
     with gr.Row():
-        output_image = gr.Gallery(label="Generated Image(s)", columns=2)
-        output_video = gr.Video(label="Generated Video")
+
+        output_image = gr.Gallery(
+            label="Generated Images",
+            columns=2
+        )
+
+        output_video = gr.Video(
+            label="Generated Video"
+        )
+
 
     user_input.submit(
-        codette_terminal_limited,
+
+        codette_terminal,
+
         inputs=[
-            user_input, model_dropdown, generate_image_toggle, generate_video_toggle,
-            session_id, batch_size_slider, video_steps_slider, fps_slider
+            user_input,
+            model_dropdown,
+            generate_image_toggle,
+            generate_video_toggle,
+            session_id,
+            batch_size_slider,
+            video_steps_slider,
+            fps_slider
         ],
-        outputs=[output_text, output_image, output_video]
+
+        outputs=[
+            output_text,
+            output_image,
+            output_video
+        ]
+
     )
 
-# ---------- Launch ----------
+
+# -------------------------------------------------
+# Launch
+# -------------------------------------------------
+
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch()