Update app.py

app.py

@@ -5,6 +5,7 @@ import streamlit as st
 import csv
 import time
 from dataclasses import dataclass
+import zipfile
 
 st.set_page_config(page_title="SFT Tiny Titans 🚀", page_icon="🤖", layout="wide", initial_sidebar_state="expanded")
 
@@ -41,6 +42,37 @@ class ModelBuilder:
             self.tokenizer.pad_token = self.tokenizer.eos_token
         self.config = config
         self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    def fine_tune(self, csv_path):
+        from torch.utils.data import Dataset, DataLoader
+        import torch
+        class SFTDataset(Dataset):
+            def __init__(self, data, tokenizer):
+                self.data = data
+                self.tokenizer = tokenizer
+            def __len__(self):
+                return len(self.data)
+            def __getitem__(self, idx):
+                prompt = self.data[idx]["prompt"]
+                response = self.data[idx]["response"]
+                inputs = self.tokenizer(f"{prompt} {response}", return_tensors="pt", padding="max_length", max_length=128, truncation=True)
+                labels = inputs["input_ids"].clone()
+                labels[0, :len(self.tokenizer(prompt)["input_ids"][0])] = -100
+                return {"input_ids": inputs["input_ids"][0], "attention_mask": inputs["attention_mask"][0], "labels": labels[0]}
+        data = []
+        with open(csv_path, "r") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                data.append({"prompt": row["prompt"], "response": row["response"]})
+        dataset = SFTDataset(data, self.tokenizer)
+        dataloader = DataLoader(dataset, batch_size=2)
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
+        self.model.train()
+        for _ in range(1):
+            for batch in dataloader:
+                optimizer.zero_grad()
+                outputs = self.model(**{k: v.to(self.model.device) for k, v in batch.items()})
+                outputs.loss.backward()
+                optimizer.step()
     def evaluate(self, prompt: str):
         import torch
         self.model.eval()
@@ -59,6 +91,25 @@ class DiffusionBuilder:
         self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
         self.pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
         self.config = config
+    def fine_tune(self, images, texts):
+        import torch
+        from PIL import Image
+        import numpy as np
+        optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
+        self.pipeline.unet.train()
+        for _ in range(1):
+            for img, text in zip(images, texts):
+                optimizer.zero_grad()
+                img_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(self.pipeline.device)
+                latents = self.pipeline.vae.encode(img_tensor).latent_dist.sample()
+                noise = torch.randn_like(latents)
+                timesteps = torch.randint(0, 1000, (1,), device=latents.device)
+                noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
+                text_emb = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
+                pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_emb).sample
+                loss = torch.nn.functional.mse_loss(pred_noise, noise)
+                loss.backward()
+                optimizer.step()
     def generate(self, prompt: str):
         return self.pipeline(prompt, num_inference_steps=20).images[0]
 
@@ -69,18 +120,23 @@ def get_download_link(file_path, mime_type="text/plain", label="Download"):
     b64 = base64.b64encode(data).decode()
     return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
 
-def generate_filename(text_line):
+def generate_filename(sequence):
     from datetime import datetime
     import pytz
     central = pytz.timezone('US/Central')
-    timestamp = datetime.now(central).strftime("%Y%m%d_%I%M%S_%p")
-    safe_text = ''.join(c if c.isalnum() else '_' for c in text_line[:50])
-    return f"{timestamp}_{safe_text}.png"
+    timestamp = datetime.now(central).strftime("%d%m%Y%H%M%S%p")
+    return f"{sequence}{timestamp}.png"
 
 def get_gallery_files(file_types):
     import glob
     return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
 
+def zip_files(files, zip_name):
+    with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for file in files:
+            zipf.write(file, os.path.basename(file))
+    return zip_name
+
 # Video Processor for WebRTC
 class VideoSnapshot:
     def __init__(self):
@@ -94,28 +150,26 @@ class VideoSnapshot:
         return self.snapshot
 
 # Main App
-st.title("SFT Tiny Titans 🚀 (Fast & Fixed!)")
+st.title("SFT Tiny Titans 🚀 (Capture & Tune!)")
 
 # Sidebar Galleries
-st.sidebar.header("Media Gallery 🎨")
-for gallery_type, file_types, emoji in [("Images 📸", ["png", "jpg", "jpeg"], "🖼️"), ("Videos 🎥", ["mp4"], "🎬")]:
-    st.sidebar.subheader(f"{gallery_type} {emoji}")
-    files = get_gallery_files(file_types)
-    if files:
-        cols = st.sidebar.columns(2)
-        for idx, file in enumerate(files[:4]):
-            with cols[idx % 2]:
-                if "Images" in gallery_type:
-                    from PIL import Image
-                    st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)
-                elif "Videos" in gallery_type:
-                    st.video(file)
+st.sidebar.header("Captured Images 🎨")
+image_files = get_gallery_files(["png"])
+if image_files:
+    cols = st.sidebar.columns(2)
+    for idx, file in enumerate(image_files[:4]):
+        with cols[idx % 2]:
+            from PIL import Image
+            st.image(Image.open(file), caption=file.split('/')[-1], use_container_width=True)
 
 # Sidebar Model Management
 st.sidebar.subheader("Model Hub 🗂️")
 model_type = st.sidebar.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"])
-model_options = {"NLP (Causal LM)": "HuggingFaceTB/SmolLM-135M", "CV (Diffusion)": "CompVis/stable-diffusion-v1-4"}
-selected_model = st.sidebar.selectbox("Select Model", ["None", model_options[model_type]])
+model_options = {
+    "NLP (Causal LM)": "HuggingFaceTB/SmolLM-135M",
+    "CV (Diffusion)": ["CompVis/stable-diffusion-v1-4", "stabilityai/stable-diffusion-2-base", "runwayml/stable-diffusion-v1-5"]
+}
+selected_model = st.sidebar.selectbox("Select Model", ["None"] + ([model_options[model_type]] if "NLP" in model_type else model_options[model_type]))
 if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
     config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=selected_model)
@@ -130,7 +184,7 @@ tab1, tab2, tab3, tab4 = st.tabs(["Build Titan 🌱", "Fine-Tune Titans 🔧", "
 with tab1:
     st.header("Build Titan 🌱 (Quick Start!)")
     model_type = st.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"], key="build_type")
-    base_model = st.selectbox("Select Model", [model_options[model_type]], key="build_model")
+    base_model = st.selectbox("Select Model", model_options[model_type], key="build_model")
     if st.button("Download Model ⬇️"):
         config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=f"titan_{int(time.time())}", base_model=base_model)
         builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
@@ -149,67 +203,22 @@ with tab2:
             st.subheader("NLP Tune 🧠")
             uploaded_csv = st.file_uploader("Upload CSV", type="csv", key="nlp_csv")
             if uploaded_csv and st.button("Tune NLP 🔄"):
-                from torch.utils.data import Dataset, DataLoader
-                import torch
-                class SFTDataset(Dataset):
-                    def __init__(self, data, tokenizer):
-                        self.data = data
-                        self.tokenizer = tokenizer
-                    def __len__(self):
-                        return len(self.data)
-                    def __getitem__(self, idx):
-                        prompt = self.data[idx]["prompt"]
-                        response = self.data[idx]["response"]
-                        inputs = self.tokenizer(f"{prompt} {response}", return_tensors="pt", padding="max_length", max_length=128, truncation=True)
-                        labels = inputs["input_ids"].clone()
-                        labels[0, :len(self.tokenizer(prompt)["input_ids"][0])] = -100
-                        return {"input_ids": inputs["input_ids"][0], "attention_mask": inputs["attention_mask"][0], "labels": labels[0]}
-                data = []
                 with open("temp.csv", "wb") as f:
                     f.write(uploaded_csv.read())
-                with open("temp.csv", "r") as f:
-                    reader = csv.DictReader(f)
-                    for row in reader:
-                        data.append({"prompt": row["prompt"], "response": row["response"]})
-                dataset = SFTDataset(data, st.session_state['builder'].tokenizer)
-                dataloader = DataLoader(dataset, batch_size=2)
-                optimizer = torch.optim.AdamW(st.session_state['builder'].model.parameters(), lr=2e-5)
-                st.session_state['builder'].model.train()
-                for _ in range(1):
-                    for batch in dataloader:
-                        optimizer.zero_grad()
-                        outputs = st.session_state['builder'].model(**{k: v.to(st.session_state['builder'].model.device) for k, v in batch.items()})
-                        outputs.loss.backward()
-                        optimizer.step()
+                st.session_state['builder'].fine_tune("temp.csv")
                 st.success("NLP sharpened! 🎉")
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
             st.subheader("CV Tune 🎨")
-            uploaded_files = st.file_uploader("Upload Images", type=["png", "jpg"], accept_multiple_files=True, key="cv_upload")
-            text_input = st.text_area("Text (one per image)", "Bat Neon\nIron Glow", key="cv_text")
-            if uploaded_files and st.button("Tune CV 🔄"):
-                import torch
-                from PIL import Image
-                import numpy as np
-                images = [Image.open(f).convert("RGB") for f in uploaded_files]
-                texts = text_input.splitlines()[:len(images)]
-                optimizer = torch.optim.AdamW(st.session_state['builder'].pipeline.unet.parameters(), lr=1e-5)
-                st.session_state['builder'].pipeline.unet.train()
-                for _ in range(1):
-                    for img, text in zip(images, texts):
-                        optimizer.zero_grad()
-                        latents = st.session_state['builder'].pipeline.vae.encode(torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float().to(st.session_state['builder'].pipeline.device)).latent_dist.sample()
-                        noise = torch.randn_like(latents)
-                        timesteps = torch.randint(0, 1000, (1,), device=latents.device)
-                        noisy_latents = st.session_state['builder'].pipeline.scheduler.add_noise(latents, noise, timesteps)
-                        text_emb = st.session_state['builder'].pipeline.text_encoder(st.session_state['builder'].pipeline.tokenizer(text, return_tensors="pt").input_ids.to(st.session_state['builder'].pipeline.device))[0]
-                        pred_noise = st.session_state['builder'].pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_emb).sample
-                        loss = torch.nn.functional.mse_loss(pred_noise, noise)
-                        loss.backward()
-                        optimizer.step()
-                for img, text in zip(images, texts):
-                    filename = generate_filename(text)
-                    img.save(filename)
-                st.success("CV polished! 🎉")
+            captured_images = get_gallery_files(["png"])
+            if len(captured_images) >= 2:
+                texts = ["Superhero Neon", "Hero Glow", "Cape Spark"][:len(captured_images)]
+                if st.button("Tune CV 🔄"):
+                    from PIL import Image
+                    images = [Image.open(img) for img in captured_images]
+                    st.session_state['builder'].fine_tune(images, texts)
+                    st.success("CV polished! 🎉")
+            else:
+                st.warning("Capture at least 2 images first! ⚠️")
 
 with tab3:
     st.header("Test Titans 🧪 (Quick Check!)")
@@ -228,10 +237,10 @@ with tab3:
             if st.button("Test CV ▶️"):
                 with st.spinner("Generating... ⏳"):
                     img = st.session_state['builder'].generate(prompt)
-                st.image(img, caption="Generated Art")
+                st.image(img, caption="Generated Art", use_container_width=True)
 
 with tab4:
-    st.header("Camera Snap 📷 (Instant Shots!)")
+    st.header("Camera Snap 📷 (Sequence Shots!)")
     from streamlit_webrtc import webrtc_streamer
     ctx = webrtc_streamer(
         key="camera",
@@ -239,29 +248,36 @@ with tab4:
         frontend_rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
     )
     if ctx.video_processor:
-        snapshot_text = st.text_input("Snapshot Text", "Live Snap")
-        if st.button("Snap It! 📸"):
-            snapshot = ctx.video_processor.take_snapshot()
-            if snapshot:
-                filename = generate_filename(snapshot_text)
-                snapshot.save(filename)
-                st.image(snapshot, caption=filename, use_container_width=True)
-                st.success("Snapped! 🎉")
+        delay = st.slider("Delay between captures (seconds)", 0, 10, 2)
+        if st.button("Capture 6 Frames 📸"):
+            captured_images = []
+            for i in range(6):
+                snapshot = ctx.video_processor.take_snapshot()
+                if snapshot:
+                    filename = generate_filename(i)
+                    snapshot.save(filename)
+                    st.image(snapshot, caption=filename, use_container_width=True)
+                    captured_images.append(filename)
+                time.sleep(delay)
+            st.success("6 frames captured! 🎉")
+            if len(captured_images) >= 2:
+                st.session_state['captured_images'] = captured_images
 
-    # Demo Dataset
-    st.subheader("Demo CV Dataset 🎨")
-    demo_texts = ["Bat Neon", "Iron Glow"]
-    demo_images = [generate_filename(t) for t in demo_texts]
-    for img, text in zip(demo_images, demo_texts):
-        if not os.path.exists(img):
-            from PIL import Image
-            Image.new("RGB", (100, 100)).save(img)
-    st.code("\n".join([f"{i+1}. {t} -> {img}" for i, (t, img) in enumerate(zip(demo_texts, demo_images))]), language="text")
-    if st.button("Download Demo CSV 📝"):
-        csv_path = f"demo_cv_{int(time.time())}.csv"
-        with open(csv_path, "w", newline="") as f:
-            writer = csv.writer(f)
-            writer.writerow(["image", "text"])
-            for img, text in zip(demo_images, demo_texts):
-                writer.writerow([img, text])
-        st.markdown(get_download_link(csv_path, "text/csv", "Download Demo CSV"), unsafe_allow_html=True)
+    # Dataset and ZIP Download
+    if 'captured_images' in st.session_state and len(st.session_state['captured_images']) >= 2:
+        st.subheader("Diffusion SFT Dataset 🎨")
+        sample_texts = ["Neon Hero", "Glowing Cape", "Spark Flyer", "Dark Knight", "Iron Shine", "Thunder Bolt"]
+        dataset = list(zip(st.session_state['captured_images'], sample_texts[:len(st.session_state['captured_images'])]))
+        st.code("\n".join([f"{i+1}. {text} -> {img}" for i, (img, text) in enumerate(dataset)]), language="text")
+        if st.button("Download Dataset CSV 📝"):
+            csv_path = f"diffusion_sft_{int(time.time())}.csv"
+            with open(csv_path, "w", newline="") as f:
+                writer = csv.writer(f)
+                writer.writerow(["image", "text"])
+                for img, text in dataset:
+                    writer.writerow([img, text])
+            st.markdown(get_download_link(csv_path, "text/csv", "Download Dataset CSV"), unsafe_allow_html=True)
+        if st.button("Download Images ZIP 📦"):
+            zip_path = f"captured_images_{int(time.time())}.zip"
+            zip_files(st.session_state['captured_images'], zip_path)
+            st.markdown(get_download_link(zip_path, "application/zip", "Download Images ZIP"), unsafe_allow_html=True)