Update app_flash.py

app_flash.py

@@ -1,227 +1,117 @@
 import os
-import re
 import torch
-import torch.nn as nn
-import torch.optim as optim
-from typing import Tuple
-from datasets import load_dataset
-from flashpack import FlashPackMixin
-from huggingface_hub import HfApi, create_repo, repo_exists
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
-
-# ============================================================
-# ⚙️ Setup
-# ============================================================
-device = torch.device("cpu")
-torch.set_num_threads(4)
-print(f"🔧 Using device: {device} (CPU-only mode)")
-
-HF_REPO = "rahul7star/FlashPack"
-MODEL_ID = HF_REPO
-
-
-# ============================================================
-# 🧠 Define FlashPack Trainer
-# ============================================================
-class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim=768, hidden_dim=512, output_dim=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):
-        return self.fc2(self.relu(self.fc1(x)))
+from diffusers import DiffusionPipeline
+from flashpack.integrations.diffusers import (
+    FlashPackDiffusersModelMixin,
+    FlashPackDiffusionPipeline,
+)
+from huggingface_hub import snapshot_download
 
 
 # ============================================================
-# 🔤 Encoder Builder (GPT2 base)
+# 🧠 Device setup (CPU fallback safe)
 # ============================================================
-def build_encoder(model_name="gpt2", max_length=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(text: str):
-        inputs = tokenizer(
-            text,
-            return_tensors="pt",
-            truncation=True,
-            padding="max_length",
-            max_length=max_length,
-        ).to(device)
-        return embed_model(**inputs).last_hidden_state.mean(dim=1).cpu()
-
-    return tokenizer, embed_model, encode
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🔧 Using device: {device}")
 
 
 # ============================================================
-# 🧩 FlashPack: Train and Upload (uses Gemma only internally)
+# 🧩 Define FlashPack-integrated pipeline
 # ============================================================
-def train_flashpack_model(hf_repo=HF_REPO):
-    print(f"🚀 Training new FlashPack model for repo: {hf_repo}")
-    model = GemmaTrainer()
-    tokenizer, embed_model, encode = build_encoder("gpt2")
-
-    # Load dataset (Gemma-expanded dataset)
-    dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
-
-    # Compute embeddings for training (short → long)
-    X, Y = [], []
-    for p in dataset.select(range(300)):
-        short_emb = encode(p["short_prompt"])
-        long_emb = encode(p["long_prompt"])
-        X.append(short_emb)
-        Y.append(long_emb)
-
-    X = torch.vstack(X)
-    Y = torch.vstack(Y)
-
-    optimizer = optim.Adam(model.parameters(), lr=1e-3)
-    for epoch in range(10):
-        out = model(X)
-        loss = nn.MSELoss()(out, Y)
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        print(f"Epoch {epoch+1}/10 | Loss: {loss.item():.6f}")
-
-    # Save FlashPack model and push
-    model.to_flashpack("flashpack_model")
-    print("💾 Model saved locally. Uploading to Hugging Face...")
-
-    api = HfApi()
-    if not repo_exists(hf_repo):
-        create_repo(hf_repo, repo_type="model", exist_ok=True)
-    model.push_to_hub(hf_repo, commit_message="Initial FlashPack model training")
-
-    print(f"✅ Model uploaded successfully to {hf_repo}")
-    return model, tokenizer, embed_model
+class FlashPackMyPipeline(DiffusionPipeline, FlashPackDiffusionPipeline):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
 
 # ============================================================
-# 📦 Load FlashPack from Hub
+# 🚀 Load FlashPack pipeline
 # ============================================================
-from huggingface_hub import snapshot_download
-import os
-
-def load_flashpack_model(hf_repo="rahul7star/FlashPack"):
-    print(f"🔁 Loading FlashPack model from: {hf_repo}")
-
-    # Try local first, then Hugging Face Hub
-    if os.path.isdir(hf_repo):
-        local_dir = hf_repo
-        print(f"📂 Using local FlashPack model at: {local_dir}")
-    else:
-        print("☁️ Downloading FlashPack model from Hugging Face Hub...")
-        local_dir = snapshot_download(repo_id=hf_repo)
-        print(f"📥 Model snapshot downloaded to: {local_dir}")
-
-    # Load from local directory
-    model = GemmaTrainer.from_flashpack(local_dir)
-    model.eval()
-    print("✅ FlashPack model loaded successfully.")
-    return model
+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}")
 
+    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}")
 
+    pipeline.to(device)
+    return pipeline
 
 
 # ============================================================
-# ⚡ Auto Load or Train
+# 🧪 Inference function
 # ============================================================
-def get_flashpack_model(hf_repo=HF_REPO):
+def generate_from_prompt(prompt: str):
+    if not prompt or prompt.strip() == "":
+        return "Please enter a valid prompt.", None
+
     try:
-        api = HfApi()
-        if repo_exists(hf_repo):
-            print("✅ Found trained model on Hub.")
-            return load_flashpack_model(hf_repo)
+        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:
-            print("❌ Model not found, training new one using Gemma dataset...")
-            return train_flashpack_model(hf_repo)
+            return "⚠️ Unknown output format.", None
     except Exception as e:
-        print(f"⚠️ Repo check failed: {e}. Retraining model locally.")
-        return train_flashpack_model(hf_repo)
+        return f"❌ Inference error: {str(e)}", None
 
 
 # ============================================================
-# 📚 Dataset + Model
+# ⚙️ Load the model
 # ============================================================
-model, tokenizer, embed_model = get_flashpack_model()
-dataset = load_dataset("gokaygokay/prompt-enhancer-dataset", split="train")
-
-long_embeddings = torch.vstack([
-    embed_model(**tokenizer(
-        p["long_prompt"], return_tensors="pt",
-        truncation=True, padding="max_length", max_length=32
-    )).last_hidden_state.mean(dim=1).cpu()
-    for p in dataset.select(range(min(500, len(dataset))))
-])
-
-print("✅ FlashPack model and embeddings loaded.")
+try:
+    pipeline = load_flashpack_pipeline("rahul7star/FlashPack")
+except Exception as e:
+    raise SystemExit(f"🚫 Failed to load model: {e}")
 
 
 # ============================================================
-# 🧠 Inference Helpers
+# 🧠 Gradio UI
 # ============================================================
-@torch.no_grad()
-def encode_for_inference(prompt: str):
-    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()
-
-
-@torch.no_grad()
-def enhance_prompt_flashpack(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)
-    sims /= (long_embeddings.norm(dim=1) * (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
-# ============================================================
-with gr.Blocks(title="Prompt Enhancer – FlashPack Only", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="FlashPack Model – rahul7star/FlashPack", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # ✨ FlashPack Prompt Enhancer  
-    - Uses pre-trained **FlashPack model** (`rahul7star/FlashPack`)  
-    - Matches short prompts to enhanced long prompts using learned embeddings  
-    - CPU-only, no Gemma dependency during inference.
+    # ⚡ FlashPack Model Inference
+    - Loaded from **rahul7star/FlashPack**
+    - Supports both image and video outputs (depending on model type)
     """)
 
     with gr.Row():
-        chatbot = gr.Chatbot(height=420, label="Enhanced Prompts", type="messages")
         with gr.Column(scale=1):
-            user_prompt = gr.Textbox(placeholder="Enter a short prompt...", label="Your Prompt", lines=3)
-            temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
-            max_tokens = gr.Slider(32, 512, value=256, label="Max Tokens")
-            send_flashpack = gr.Button("🔗 Enhance Prompt")
-            clear_btn = gr.Button("🧹 Clear Chat")
-
-    send_flashpack.click(enhance_prompt_flashpack, [user_prompt, temperature, max_tokens, chatbot], chatbot)
-    user_prompt.submit(enhance_prompt_flashpack, [user_prompt, temperature, max_tokens, chatbot], chatbot)
-    clear_btn.click(lambda: [], None, chatbot)
+            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],
+    )
 
 # ============================================================
-# 🚀 Launch App
+# 🏁 Launch app
 # ============================================================
 if __name__ == "__main__":
     demo.launch(show_error=True)