Update app_flash.py

app_flash.py

@@ -1,114 +1,235 @@
+import gc
 import os
 import torch
+import torch.nn as nn
+import torch.optim as optim
+import tempfile
 import gradio as gr
-from diffusers import DiffusionPipeline
-from flashpack.integrations.diffusers import (
-    FlashPackDiffusersModelMixin,
-    FlashPackDiffusionPipeline,
-)
-from huggingface_hub import snapshot_download
-
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModel
+from flashpack import FlashPackMixin
+from huggingface_hub import Repository
+from typing import Tuple
 
 # ============================================================
-# 🧠 Device setup (CPU fallback safe)
+# 🖥 Device setup (CPU-only safe)
 # ============================================================
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"🔧 Using device: {device}")
-
+device = torch.device("cpu")
+torch.set_num_threads(4)
+print(f"🔧 Using device: {device}  (CPU-only mode)")
 
 # ============================================================
-# 🧩 Define FlashPack-integrated pipeline
+# 1️⃣ Define FlashPack model
 # ============================================================
-class FlashPackMyPipeline(DiffusionPipeline, FlashPackDiffusionPipeline):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+class GemmaTrainer(nn.Module, FlashPackMixin):
+    def __init__(self, input_dim: int = 768, hidden_dim: int = 512, output_dim: int = 768):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
 
+# ============================================================
+# 2️⃣ Build tokenizer + encoder
+# ============================================================
+def build_encoder(model_name="gpt2", max_length: int = 32):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    embed_model = AutoModel.from_pretrained(model_name).to(device)
+    embed_model.eval()
+
+    @torch.no_grad()
+    def encode(prompt: str) -> torch.Tensor:
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=max_length
+        ).to(device)
+        outputs = embed_model(**inputs).last_hidden_state.mean(dim=1)
+        return outputs.cpu()
+
+    return tokenizer, embed_model, encode
 
 # ============================================================
-# 🚀 Load FlashPack pipeline
+# 3️⃣ Push FlashPack model to HF
 # ============================================================
-def load_flashpack_pipeline(repo_id: str = "rahul7star/FlashPack"):
-    """
-    Loads a FlashPack pipeline from Hugging Face Hub.
-    Falls back to local snapshot if network or metadata issue occurs.
-    """
-    print(f"🔁 Loading FlashPack pipeline from: {repo_id}")
+def push_flashpack_model_to_hf(model, hf_repo: str):
+    logs = []
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        logs.append(f"📂 Using temporary directory: {tmp_dir}")
+        repo = Repository(local_dir=tmp_dir, clone_from=hf_repo, use_auth_token=True)
+        logs.append(f"🌐 Hugging Face repo cloned to: {tmp_dir}")
 
-    try:
-        # Try direct hub load
-        pipeline = FlashPackMyPipeline.from_pretrained_flashpack(repo_id)
-        print("✅ Successfully loaded FlashPack pipeline from Hugging Face Hub.")
-    except Exception as e:
-        print(f"⚠️ Hub load failed: {e}")
-        print("⏬ Attempting to load via snapshot_download...")
-        try:
-            local_dir = snapshot_download(repo_id=repo_id)
-            pipeline = FlashPackMyPipeline.from_pretrained_flashpack(local_dir)
-            print(f"✅ Loaded FlashPack pipeline from local snapshot: {local_dir}")
-        except Exception as e2:
-            raise RuntimeError(f"❌ Failed to load FlashPack model: {e2}")
+        pack_path = os.path.join(tmp_dir, "model.flashpack")
+        logs.append(f"💾 Saving model to: {pack_path}")
+        model.save_flashpack(pack_path, target_dtype=torch.float32)
+        logs.append("✅ Model saved successfully.")
 
-    pipeline.to(device)
-    return pipeline
+        readme_path = os.path.join(tmp_dir, "README.md")
+        with open(readme_path, "w") as f:
+            f.write("# FlashPack Model\nThis repo contains a FlashPack model.")
+        logs.append("📄 README.md added.")
 
+        logs.append("🚀 Pushing repo to Hugging Face Hub...")
+        repo.push_to_hub()
+        logs.append(f"✅ Model successfully pushed to: {hf_repo}")
+
+    return logs
 
 # ============================================================
-# 🧪 Inference function
+# 4️⃣ Train FlashPack model
 # ============================================================
-def generate_from_prompt(prompt: str):
-    if not prompt or prompt.strip() == "":
-        return "Please enter a valid prompt.", None
+def train_flashpack_model(
+    dataset_name: str = "gokaygokay/prompt-enhancer-dataset",
+    max_encode: int = 1000,
+    device: str = "cpu"
+) -> Tuple[GemmaTrainer, object, object, object, torch.Tensor]:
+    print("📦 Loading dataset...")
+    dataset = load_dataset(dataset_name, split="train")
+    limit = min(max_encode, len(dataset))
+    dataset = dataset.select(range(limit))
+    print(f"⚡ Encoding {len(dataset)} prompts (max {max_encode})")
+
+    tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
+
+    short_list, long_list = [], []
+    for i, item in enumerate(dataset):
+        short_list.append(encode_fn(item["short_prompt"]))
+        long_list.append(encode_fn(item["long_prompt"]))
+        if (i+1) % 50 == 0 or (i+1) == len(dataset):
+            print(f"  → Encoded {i+1}/{limit} prompts")
+            gc.collect()
+
+    short_embeddings = torch.vstack(short_list)
+    long_embeddings = torch.vstack(long_list)
+    print(f"✅ Finished encoding {short_embeddings.shape[0]} prompts")
+
+    # Build model
+    model = GemmaTrainer(
+        input_dim=short_embeddings.shape[1],
+        hidden_dim=min(512, short_embeddings.shape[1]),
+        output_dim=long_embeddings.shape[1]
+    ).to(device)
+
+    criterion = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=1e-3)
+    max_epochs = 20
+    batch_size = 32
+
+    print("🚀 Training model...")
+    n = short_embeddings.shape[0]
+    for epoch in range(max_epochs):
+        model.train()
+        epoch_loss = 0.0
+        perm = torch.randperm(n)
+        for start in range(0, n, batch_size):
+            idx = perm[start:start+batch_size]
+            inputs = short_embeddings[idx].to(device)
+            targets = long_embeddings[idx].to(device)
+
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * inputs.size(0)
+
+        epoch_loss /= n
+        if epoch % 5 == 0 or epoch == max_epochs-1:
+            print(f"Epoch {epoch+1}/{max_epochs}, Loss={epoch_loss:.6f}")
+
+    print("✅ Training finished!")
+    return model, dataset, embed_model, tokenizer, long_embeddings
 
+# ============================================================
+# 5️⃣ Load FlashPack model (train if missing)
+# ============================================================
+def get_flashpack_model(hf_repo="rahul7star/FlashPack"):
     try:
-        output = pipeline(prompt)
-        if hasattr(output, "images"):
-            img = output.images[0]
-            return f"✅ Generated successfully!", img
-        elif hasattr(output, "frames"):
-            frames = output.frames
-            video_path = "/tmp/generated.mp4"
-            from diffusers.utils import export_to_video
-            export_to_video(frames, video_path)
-            return f"✅ Video generated successfully!", video_path
-        else:
-            return "⚠️ Unknown output format.", None
+        print(f"🔁 Attempting to load FlashPack model from {hf_repo}")
+        model = GemmaTrainer.from_flashpack(hf_repo)
+        model.eval()
+        print("✅ Loaded model successfully from HF")
+        tokenizer, embed_model, encode_fn = build_encoder("gpt2", max_length=32)
+        return model, tokenizer, embed_model
     except Exception as e:
-        return f"❌ Inference error: {str(e)}", None
-
+        print(f"⚠️ Load failed: {e}")
+        print("⏬ Training a new FlashPack model locally...")
+        model, dataset, embed_model, tokenizer, long_embeddings = train_flashpack_model()
+        print("📤 Pushing trained model to HF...")
+        push_flashpack_model_to_hf(model, hf_repo)
+        return model, tokenizer, embed_model, dataset, long_embeddings
 
 # ============================================================
-# ⚙️ Load the model
+# 6️⃣ Load or train
 # ============================================================
 try:
-    pipeline = load_flashpack_pipeline("rahul7star/FlashPack")
+    model, tokenizer, embed_model, dataset, long_embeddings = get_flashpack_model()
 except Exception as e:
-    raise SystemExit(f"🚫 Failed to load model: {e}")
+    raise SystemExit(f"❌ Failed to load or train FlashPack model: {e}")
+
+# ============================================================
+# 7️⃣ Inference helpers
+# ============================================================
+@torch.no_grad()
+def encode_for_inference(prompt: str) -> torch.Tensor:
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
+                       padding="max_length", max_length=32).to(device)
+    return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
+
+def enhance_prompt(user_prompt: str, temperature: float, max_tokens: int, chat_history):
+    chat_history = chat_history or []
+    short_emb = encode_for_inference(user_prompt)
+    mapped = model(short_emb.to(device)).cpu()
 
+    sims = (long_embeddings @ mapped.t()).squeeze(1)
+    long_norms = long_embeddings.norm(dim=1)
+    mapped_norm = mapped.norm()
+    sims = sims / (long_norms * (mapped_norm + 1e-12))
+
+    best_idx = int(sims.argmax().item())
+    enhanced_prompt = dataset[best_idx]["long_prompt"]
+
+    chat_history.append({"role": "user", "content": user_prompt})
+    chat_history.append({"role": "assistant", "content": enhanced_prompt})
+    return chat_history
 
 # ============================================================
-# 🧠 Gradio UI
+# 8️⃣ Gradio UI
 # ============================================================
-with gr.Blocks(title="FlashPack Model – rahul7star/FlashPack", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # ⚡ FlashPack Model Inference
-    - Loaded from **rahul7star/FlashPack**
-    - Supports both image and video outputs (depending on model type)
-    """)
+with gr.Blocks(title="Prompt Enhancer – FlashPack (CPU)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # ✨ Prompt Enhancer (FlashPack mapper)
+        Enter a short prompt, and the model will **expand it with details and creative context**.
+        (CPU-only mode.)
+        """
+    )
 
     with gr.Row():
+        chatbot = gr.Chatbot(height=400, label="Enhanced Prompts", type="messages")
         with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g. A robot painting in the rain")
-            run_btn = gr.Button("🚀 Generate", variant="primary")
-        with gr.Column(scale=1):
-            result_msg = gr.Textbox(label="Status", interactive=False)
-            image_out = gr.Image(label="Generated Image")
-            video_out = gr.Video(label="Generated Video")
-
-    run_btn.click(
-        generate_from_prompt,
-        inputs=[prompt],
-        outputs=[result_msg, image_out],
-    )
+            user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
+            temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature")
+            max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max Tokens")
+            send_btn = gr.Button("🚀 Enhance Prompt", variant="primary")
+            clear_btn = gr.Button("🧹 Clear Chat")
+
+    send_btn.click(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    user_prompt.submit(enhance_prompt, [user_prompt, temperature, max_tokens, chatbot], chatbot)
+    clear_btn.click(lambda: [], None, chatbot)
+
+# ============================================================
+# 9️⃣ Launch
 
 # ============================================================
 # 🏁 Launch app