Update app.py
Browse files
app.py
CHANGED
@@ -19,27 +19,18 @@ import random
|
|
19 |
import logging
|
20 |
from datetime import datetime
|
21 |
import pytz
|
22 |
-
from diffusers import StableDiffusionPipeline
|
23 |
from urllib.parse import quote
|
|
|
24 |
|
25 |
-
#
|
26 |
logging.basicConfig(level=logging.INFO)
|
27 |
logger = logging.getLogger(__name__)
|
28 |
|
29 |
# Page Configuration
|
30 |
-
st.set_page_config(
|
31 |
-
page_title="SFT Tiny Titans π",
|
32 |
-
page_icon="π€",
|
33 |
-
layout="wide",
|
34 |
-
initial_sidebar_state="expanded",
|
35 |
-
menu_items={
|
36 |
-
'Get Help': 'https://huggingface.co/awacke1',
|
37 |
-
'Report a bug': 'https://huggingface.co/spaces/awacke1',
|
38 |
-
'About': "Tiny Titans: Small models, big dreams, and a sprinkle of chaos! π"
|
39 |
-
}
|
40 |
-
)
|
41 |
|
42 |
-
# Model
|
43 |
@dataclass
|
44 |
class ModelConfig:
|
45 |
name: str
|
@@ -47,7 +38,6 @@ class ModelConfig:
|
|
47 |
size: str
|
48 |
domain: Optional[str] = None
|
49 |
model_type: str = "causal_lm"
|
50 |
-
|
51 |
@property
|
52 |
def model_path(self):
|
53 |
return f"models/{self.name}"
|
@@ -57,7 +47,6 @@ class DiffusionConfig:
|
|
57 |
name: str
|
58 |
base_model: str
|
59 |
size: str
|
60 |
-
|
61 |
@property
|
62 |
def model_path(self):
|
63 |
return f"diffusion_models/{self.name}"
|
@@ -68,10 +57,8 @@ class SFTDataset(Dataset):
|
|
68 |
self.data = data
|
69 |
self.tokenizer = tokenizer
|
70 |
self.max_length = max_length
|
71 |
-
|
72 |
def __len__(self):
|
73 |
return len(self.data)
|
74 |
-
|
75 |
def __getitem__(self, idx):
|
76 |
prompt = self.data[idx]["prompt"]
|
77 |
response = self.data[idx]["response"]
|
@@ -90,133 +77,103 @@ class DiffusionDataset(Dataset):
|
|
90 |
def __init__(self, images, texts):
|
91 |
self.images = images
|
92 |
self.texts = texts
|
93 |
-
|
94 |
def __len__(self):
|
95 |
return len(self.images)
|
96 |
-
|
97 |
def __getitem__(self, idx):
|
98 |
return {"image": self.images[idx], "text": self.texts[idx]}
|
99 |
|
100 |
-
# Model
|
101 |
class ModelBuilder:
|
102 |
def __init__(self):
|
103 |
self.config = None
|
104 |
self.model = None
|
105 |
self.tokenizer = None
|
106 |
self.sft_data = None
|
107 |
-
self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! π", "Training complete! Time for a binary coffee break. β"]
|
108 |
-
|
109 |
def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
self.config = config
|
117 |
-
st.success(f"Model loaded! π {random.choice(self.jokes)}")
|
118 |
return self
|
119 |
-
|
120 |
def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
|
121 |
self.sft_data = []
|
122 |
with open(csv_path, "r") as f:
|
123 |
reader = csv.DictReader(f)
|
124 |
for row in reader:
|
125 |
self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
|
126 |
-
|
127 |
dataset = SFTDataset(self.sft_data, self.tokenizer)
|
128 |
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
|
129 |
optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
|
130 |
-
|
131 |
self.model.train()
|
132 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
133 |
self.model.to(device)
|
134 |
for epoch in range(epochs):
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
|
148 |
-
st.success(f"SFT Fine-tuning completed! π {random.choice(self.jokes)}")
|
149 |
return self
|
150 |
-
|
151 |
def save_model(self, path: str):
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
st.success(f"Model saved at {path}! β
")
|
157 |
-
|
158 |
-
def evaluate(self, prompt: str, status_container=None):
|
159 |
self.model.eval()
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
|
165 |
-
outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
|
166 |
-
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
167 |
-
except Exception as e:
|
168 |
-
if status_container:
|
169 |
-
status_container.error(f"Oops! Something broke: {str(e)} π₯")
|
170 |
-
return f"Error: {str(e)}"
|
171 |
|
172 |
class DiffusionBuilder:
|
173 |
def __init__(self):
|
174 |
self.config = None
|
175 |
self.pipeline = None
|
176 |
-
|
177 |
def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
self.config = config
|
183 |
-
st.success(f"Diffusion model loaded! π¨")
|
184 |
return self
|
185 |
-
|
186 |
def fine_tune_sft(self, images, texts, epochs=3):
|
187 |
dataset = DiffusionDataset(images, texts)
|
188 |
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
|
189 |
optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
|
190 |
-
|
191 |
self.pipeline.unet.train()
|
192 |
for epoch in range(epochs):
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
|
210 |
-
st.success("Diffusion SFT Fine-tuning completed! π¨")
|
211 |
return self
|
212 |
-
|
213 |
def save_model(self, path: str):
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
|
219 |
-
#
|
220 |
def get_download_link(file_path, mime_type="text/plain", label="Download"):
|
221 |
with open(file_path, 'rb') as f:
|
222 |
data = f.read()
|
@@ -227,19 +184,14 @@ def zip_directory(directory_path, zip_path):
|
|
227 |
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
|
228 |
for root, _, files in os.walk(directory_path):
|
229 |
for file in files:
|
230 |
-
|
231 |
-
arcname = os.path.relpath(file_path, os.path.dirname(directory_path))
|
232 |
-
zipf.write(file_path, arcname)
|
233 |
|
234 |
def get_model_files(model_type="causal_lm"):
|
235 |
path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
|
236 |
return [d for d in glob.glob(path) if os.path.isdir(d)]
|
237 |
|
238 |
def get_gallery_files(file_types):
|
239 |
-
|
240 |
-
for ext in file_types:
|
241 |
-
files.extend(glob.glob(f"*.{ext}"))
|
242 |
-
return sorted(files)
|
243 |
|
244 |
def generate_filename(text_line):
|
245 |
central = pytz.timezone('US/Central')
|
@@ -254,39 +206,55 @@ def display_search_links(query):
|
|
254 |
"Google": f"https://www.google.com/search?q={quote(query)}",
|
255 |
"YouTube": f"https://www.youtube.com/results?search_query={quote(query)}"
|
256 |
}
|
257 |
-
|
258 |
-
|
259 |
-
|
260 |
-
|
261 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
262 |
def __init__(self, model, tokenizer):
|
263 |
self.model = model
|
264 |
self.tokenizer = tokenizer
|
265 |
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
266 |
self.model.to(self.device)
|
267 |
-
|
268 |
def generate(self, prompt: str) -> str:
|
269 |
self.model.eval()
|
270 |
with torch.no_grad():
|
271 |
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device)
|
272 |
outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7)
|
273 |
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
274 |
-
|
275 |
def plan_party(self, task: str) -> pd.DataFrame:
|
276 |
search_result = "Latest trends for 2025: Gold-plated Batman statues, VR superhero battles."
|
277 |
prompt = f"Given this context: '{search_result}'\n{task}"
|
278 |
plan_text = self.generate(prompt)
|
279 |
st.markdown(f"Search Links: {display_search_links('superhero party trends')}", unsafe_allow_html=True)
|
280 |
-
|
281 |
-
|
282 |
-
|
283 |
-
|
284 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
285 |
data = [
|
286 |
-
{"
|
287 |
-
{"
|
288 |
-
{"Location": "London", "Travel Time (hrs)": travel_times["London"], "Luxury Idea": "Live stunt shows"},
|
289 |
-
{"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Luxury Idea": "Holographic displays"}
|
290 |
]
|
291 |
return pd.DataFrame(data)
|
292 |
|
@@ -309,52 +277,59 @@ def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_
|
|
309 |
st.title("SFT Tiny Titans π (Small but Mighty!)")
|
310 |
|
311 |
# Sidebar Galleries
|
312 |
-
st.sidebar.header("Galleries π¨")
|
313 |
-
for gallery_type, file_types in [
|
314 |
-
("
|
315 |
-
("
|
316 |
-
("Audio
|
317 |
]:
|
318 |
-
st.sidebar.subheader(gallery_type)
|
319 |
files = get_gallery_files(file_types)
|
320 |
if files:
|
321 |
cols_num = st.sidebar.slider(f"{gallery_type} Columns", 1, 5, 3, key=f"{gallery_type}_cols")
|
322 |
cols = st.sidebar.columns(cols_num)
|
323 |
for idx, file in enumerate(files[:cols_num * 2]):
|
324 |
with cols[idx % cols_num]:
|
325 |
-
if "
|
326 |
st.image(Image.open(file), caption=file, use_column_width=True)
|
327 |
-
elif "
|
328 |
st.video(file)
|
329 |
elif "Audio" in gallery_type:
|
330 |
st.audio(file)
|
331 |
|
332 |
st.sidebar.subheader("Model Management ποΈ")
|
333 |
-
model_type = st.sidebar.selectbox("Model Type", ["Causal LM", "Diffusion"])
|
334 |
-
model_dirs = get_model_files("causal_lm" if
|
335 |
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs)
|
336 |
if selected_model != "None" and st.sidebar.button("Load Model π"):
|
337 |
-
if
|
338 |
-
|
339 |
-
|
340 |
-
st.session_state['builder']
|
341 |
st.session_state['model_loaded'] = True
|
342 |
st.rerun()
|
343 |
|
344 |
# Tabs
|
345 |
-
tab1, tab2, tab3, tab4, tab5 = st.tabs([
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
346 |
|
347 |
with tab1:
|
348 |
-
st.header("Build
|
349 |
-
model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
|
350 |
-
|
351 |
-
|
352 |
-
|
353 |
-
|
354 |
model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
|
355 |
if st.button("Download Model β¬οΈ"):
|
356 |
-
config = (ModelConfig if
|
357 |
-
builder = ModelBuilder() if
|
358 |
builder.load_model(base_model, config)
|
359 |
builder.save_model(config.model_path)
|
360 |
st.session_state['builder'] = builder
|
@@ -362,72 +337,146 @@ with tab1:
|
|
362 |
st.rerun()
|
363 |
|
364 |
with tab2:
|
365 |
-
st.header("Fine-Tune Titan
|
366 |
-
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
|
367 |
-
st.warning("Please build or load a Titan first! β οΈ")
|
368 |
-
else:
|
369 |
-
if isinstance(st.session_state['builder'], ModelBuilder):
|
370 |
-
uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv")
|
371 |
-
if uploaded_csv and st.button("Fine-Tune with Uploaded CSV π"):
|
372 |
-
csv_path = f"uploaded_sft_data_{int(time.time())}.csv"
|
373 |
-
with open(csv_path, "wb") as f:
|
374 |
-
f.write(uploaded_csv.read())
|
375 |
-
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
|
376 |
-
new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
|
377 |
-
st.session_state['builder'].config = new_config
|
378 |
-
st.session_state['builder'].fine_tune_sft(csv_path)
|
379 |
-
st.session_state['builder'].save_model(new_config.model_path)
|
380 |
-
zip_path = f"{new_config.model_path}.zip"
|
381 |
-
zip_directory(new_config.model_path, zip_path)
|
382 |
-
st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
|
383 |
-
|
384 |
-
with tab3:
|
385 |
-
st.header("Test Titan π§ͺ")
|
386 |
-
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
|
387 |
-
st.warning("Please build or load a Titan first! β οΈ")
|
388 |
-
else:
|
389 |
-
if isinstance(st.session_state['builder'], ModelBuilder):
|
390 |
-
test_prompt = st.text_area("Enter Test Prompt", "What is AI?")
|
391 |
-
if st.button("Run Test βΆοΈ"):
|
392 |
-
result = st.session_state['builder'].evaluate(test_prompt)
|
393 |
-
st.write(f"**Generated Response**: {result}")
|
394 |
-
|
395 |
-
with tab4:
|
396 |
-
st.header("Agentic RAG Party π")
|
397 |
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], ModelBuilder):
|
398 |
-
st.warning("
|
399 |
else:
|
400 |
-
|
401 |
-
|
402 |
-
|
403 |
-
|
404 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
405 |
|
406 |
-
with
|
407 |
-
st.header("
|
408 |
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], DiffusionBuilder):
|
409 |
-
st.warning("
|
410 |
else:
|
411 |
-
uploaded_files = st.file_uploader("Upload Images/Videos", type=["png", "jpg", "jpeg", "mp4", "mp3"], accept_multiple_files=True)
|
412 |
-
text_input = st.text_area("Enter Text (one line per image)", "
|
413 |
-
if uploaded_files and st.button("
|
414 |
images = [Image.open(f) for f in uploaded_files if f.type.startswith("image")]
|
415 |
texts = text_input.splitlines()
|
416 |
if len(images) > len(texts):
|
417 |
texts.extend([""] * (len(images) - len(texts)))
|
418 |
elif len(texts) > len(images):
|
419 |
texts = texts[:len(images)]
|
420 |
-
|
421 |
st.session_state['builder'].fine_tune_sft(images, texts)
|
422 |
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
|
423 |
new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
|
424 |
st.session_state['builder'].config = new_config
|
425 |
st.session_state['builder'].save_model(new_config.model_path)
|
426 |
-
|
427 |
for img, text in zip(images, texts):
|
428 |
filename = generate_filename(text)
|
429 |
img.save(filename)
|
430 |
st.image(img, caption=filename)
|
431 |
zip_path = f"{new_config.model_path}.zip"
|
432 |
zip_directory(new_config.model_path, zip_path)
|
433 |
-
st.markdown(get_download_link(zip_path, "application/zip", "Download
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
19 |
import logging
|
20 |
from datetime import datetime
|
21 |
import pytz
|
22 |
+
from diffusers import StableDiffusionPipeline
|
23 |
from urllib.parse import quote
|
24 |
+
import cv2
|
25 |
|
26 |
+
# Logging setup
|
27 |
logging.basicConfig(level=logging.INFO)
|
28 |
logger = logging.getLogger(__name__)
|
29 |
|
30 |
# Page Configuration
|
31 |
+
st.set_page_config(page_title="SFT Tiny Titans π", page_icon="π€", layout="wide", initial_sidebar_state="expanded")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
32 |
|
33 |
+
# Model Configurations
|
34 |
@dataclass
|
35 |
class ModelConfig:
|
36 |
name: str
|
|
|
38 |
size: str
|
39 |
domain: Optional[str] = None
|
40 |
model_type: str = "causal_lm"
|
|
|
41 |
@property
|
42 |
def model_path(self):
|
43 |
return f"models/{self.name}"
|
|
|
47 |
name: str
|
48 |
base_model: str
|
49 |
size: str
|
|
|
50 |
@property
|
51 |
def model_path(self):
|
52 |
return f"diffusion_models/{self.name}"
|
|
|
57 |
self.data = data
|
58 |
self.tokenizer = tokenizer
|
59 |
self.max_length = max_length
|
|
|
60 |
def __len__(self):
|
61 |
return len(self.data)
|
|
|
62 |
def __getitem__(self, idx):
|
63 |
prompt = self.data[idx]["prompt"]
|
64 |
response = self.data[idx]["response"]
|
|
|
77 |
def __init__(self, images, texts):
|
78 |
self.images = images
|
79 |
self.texts = texts
|
|
|
80 |
def __len__(self):
|
81 |
return len(self.images)
|
|
|
82 |
def __getitem__(self, idx):
|
83 |
return {"image": self.images[idx], "text": self.texts[idx]}
|
84 |
|
85 |
+
# Model Builders
|
86 |
class ModelBuilder:
|
87 |
def __init__(self):
|
88 |
self.config = None
|
89 |
self.model = None
|
90 |
self.tokenizer = None
|
91 |
self.sft_data = None
|
|
|
|
|
92 |
def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
|
93 |
+
self.model = AutoModelForCausalLM.from_pretrained(model_path)
|
94 |
+
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
95 |
+
if self.tokenizer.pad_token is None:
|
96 |
+
self.tokenizer.pad_token = self.tokenizer.eos_token
|
97 |
+
if config:
|
98 |
+
self.config = config
|
|
|
|
|
99 |
return self
|
|
|
100 |
def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
|
101 |
self.sft_data = []
|
102 |
with open(csv_path, "r") as f:
|
103 |
reader = csv.DictReader(f)
|
104 |
for row in reader:
|
105 |
self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
|
|
|
106 |
dataset = SFTDataset(self.sft_data, self.tokenizer)
|
107 |
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
|
108 |
optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
|
|
|
109 |
self.model.train()
|
110 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
111 |
self.model.to(device)
|
112 |
for epoch in range(epochs):
|
113 |
+
total_loss = 0
|
114 |
+
for batch in dataloader:
|
115 |
+
optimizer.zero_grad()
|
116 |
+
input_ids = batch["input_ids"].to(device)
|
117 |
+
attention_mask = batch["attention_mask"].to(device)
|
118 |
+
labels = batch["labels"].to(device)
|
119 |
+
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
|
120 |
+
loss = outputs.loss
|
121 |
+
loss.backward()
|
122 |
+
optimizer.step()
|
123 |
+
total_loss += loss.item()
|
124 |
+
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
|
|
|
|
|
125 |
return self
|
|
|
126 |
def save_model(self, path: str):
|
127 |
+
os.makedirs(os.path.dirname(path), exist_ok=True)
|
128 |
+
self.model.save_pretrained(path)
|
129 |
+
self.tokenizer.save_pretrained(path)
|
130 |
+
def evaluate(self, prompt: str):
|
|
|
|
|
|
|
131 |
self.model.eval()
|
132 |
+
with torch.no_grad():
|
133 |
+
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
|
134 |
+
outputs = self.model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
|
135 |
+
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
136 |
|
137 |
class DiffusionBuilder:
|
138 |
def __init__(self):
|
139 |
self.config = None
|
140 |
self.pipeline = None
|
|
|
141 |
def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None):
|
142 |
+
self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
|
143 |
+
self.pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
|
144 |
+
if config:
|
145 |
+
self.config = config
|
|
|
|
|
146 |
return self
|
|
|
147 |
def fine_tune_sft(self, images, texts, epochs=3):
|
148 |
dataset = DiffusionDataset(images, texts)
|
149 |
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
|
150 |
optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
|
|
|
151 |
self.pipeline.unet.train()
|
152 |
for epoch in range(epochs):
|
153 |
+
total_loss = 0
|
154 |
+
for batch in dataloader:
|
155 |
+
optimizer.zero_grad()
|
156 |
+
image = batch["image"].to(self.pipeline.device)
|
157 |
+
text = batch["text"]
|
158 |
+
latents = self.pipeline.vae.encode(image).latent_dist.sample()
|
159 |
+
noise = torch.randn_like(latents)
|
160 |
+
timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device)
|
161 |
+
noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
|
162 |
+
text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
|
163 |
+
pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
|
164 |
+
loss = torch.nn.functional.mse_loss(pred_noise, noise)
|
165 |
+
loss.backward()
|
166 |
+
optimizer.step()
|
167 |
+
total_loss += loss.item()
|
168 |
+
st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
|
|
|
|
|
169 |
return self
|
|
|
170 |
def save_model(self, path: str):
|
171 |
+
os.makedirs(os.path.dirname(path), exist_ok=True)
|
172 |
+
self.pipeline.save_pretrained(path)
|
173 |
+
def generate(self, prompt: str):
|
174 |
+
return self.pipeline(prompt, num_inference_steps=50).images[0]
|
175 |
|
176 |
+
# Utilities
|
177 |
def get_download_link(file_path, mime_type="text/plain", label="Download"):
|
178 |
with open(file_path, 'rb') as f:
|
179 |
data = f.read()
|
|
|
184 |
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
|
185 |
for root, _, files in os.walk(directory_path):
|
186 |
for file in files:
|
187 |
+
zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.dirname(directory_path)))
|
|
|
|
|
188 |
|
189 |
def get_model_files(model_type="causal_lm"):
|
190 |
path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
|
191 |
return [d for d in glob.glob(path) if os.path.isdir(d)]
|
192 |
|
193 |
def get_gallery_files(file_types):
|
194 |
+
return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
|
|
|
|
|
|
|
195 |
|
196 |
def generate_filename(text_line):
|
197 |
central = pytz.timezone('US/Central')
|
|
|
206 |
"Google": f"https://www.google.com/search?q={quote(query)}",
|
207 |
"YouTube": f"https://www.youtube.com/results?search_query={quote(query)}"
|
208 |
}
|
209 |
+
return ' '.join([f"[{name}]({url})" for name, url in search_urls.items()])
|
210 |
+
|
211 |
+
def detect_cameras():
|
212 |
+
cameras = []
|
213 |
+
for i in range(2): # Check first two indices
|
214 |
+
cap = cv2.VideoCapture(i)
|
215 |
+
if cap.isOpened():
|
216 |
+
cameras.append(i)
|
217 |
+
cap.release()
|
218 |
+
return cameras
|
219 |
+
|
220 |
+
# Agent Classes
|
221 |
+
class NLPAgent:
|
222 |
def __init__(self, model, tokenizer):
|
223 |
self.model = model
|
224 |
self.tokenizer = tokenizer
|
225 |
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
226 |
self.model.to(self.device)
|
|
|
227 |
def generate(self, prompt: str) -> str:
|
228 |
self.model.eval()
|
229 |
with torch.no_grad():
|
230 |
inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.device)
|
231 |
outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.7)
|
232 |
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
233 |
def plan_party(self, task: str) -> pd.DataFrame:
|
234 |
search_result = "Latest trends for 2025: Gold-plated Batman statues, VR superhero battles."
|
235 |
prompt = f"Given this context: '{search_result}'\n{task}"
|
236 |
plan_text = self.generate(prompt)
|
237 |
st.markdown(f"Search Links: {display_search_links('superhero party trends')}", unsafe_allow_html=True)
|
238 |
+
locations = {"Wayne Manor": (42.3601, -71.0589), "New York": (40.7128, -74.0060)}
|
239 |
+
travel_times = {loc: calculate_cargo_travel_time(coords, locations["Wayne Manor"]) for loc, coords in locations.items() if loc != "Wayne Manor"}
|
240 |
+
data = [
|
241 |
+
{"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Idea": "Gold-plated Batman statues"},
|
242 |
+
{"Location": "Wayne Manor", "Travel Time (hrs)": 0.0, "Idea": "VR superhero battles"}
|
243 |
+
]
|
244 |
+
return pd.DataFrame(data)
|
245 |
+
|
246 |
+
class CVAgent:
|
247 |
+
def __init__(self, pipeline):
|
248 |
+
self.pipeline = pipeline
|
249 |
+
def generate(self, prompt: str) -> Image.Image:
|
250 |
+
return self.pipeline(prompt, num_inference_steps=50).images[0]
|
251 |
+
def enhance_images(self, task: str) -> pd.DataFrame:
|
252 |
+
search_result = "Latest superhero art trends: Neon outlines, 3D holograms."
|
253 |
+
prompt = f"Given this context: '{search_result}'\n{task}"
|
254 |
+
st.markdown(f"Search Links: {display_search_links('superhero art trends')}", unsafe_allow_html=True)
|
255 |
data = [
|
256 |
+
{"Image Theme": "Batman", "Enhancement": "Neon outlines"},
|
257 |
+
{"Image Theme": "Iron Man", "Enhancement": "3D holograms"}
|
|
|
|
|
258 |
]
|
259 |
return pd.DataFrame(data)
|
260 |
|
|
|
277 |
st.title("SFT Tiny Titans π (Small but Mighty!)")
|
278 |
|
279 |
# Sidebar Galleries
|
280 |
+
st.sidebar.header("Shared Galleries π¨")
|
281 |
+
for gallery_type, file_types, emoji in [
|
282 |
+
("Images πΈ", ["png", "jpg", "jpeg"], "πΌοΈ"),
|
283 |
+
("Videos π₯", ["mp4"], "π¬"),
|
284 |
+
("Audio πΆ", ["mp3"], "π΅")
|
285 |
]:
|
286 |
+
st.sidebar.subheader(f"{gallery_type} {emoji}")
|
287 |
files = get_gallery_files(file_types)
|
288 |
if files:
|
289 |
cols_num = st.sidebar.slider(f"{gallery_type} Columns", 1, 5, 3, key=f"{gallery_type}_cols")
|
290 |
cols = st.sidebar.columns(cols_num)
|
291 |
for idx, file in enumerate(files[:cols_num * 2]):
|
292 |
with cols[idx % cols_num]:
|
293 |
+
if "Images" in gallery_type:
|
294 |
st.image(Image.open(file), caption=file, use_column_width=True)
|
295 |
+
elif "Videos" in gallery_type:
|
296 |
st.video(file)
|
297 |
elif "Audio" in gallery_type:
|
298 |
st.audio(file)
|
299 |
|
300 |
st.sidebar.subheader("Model Management ποΈ")
|
301 |
+
model_type = st.sidebar.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"])
|
302 |
+
model_dirs = get_model_files("causal_lm" if "NLP" in model_type else "diffusion")
|
303 |
selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs)
|
304 |
if selected_model != "None" and st.sidebar.button("Load Model π"):
|
305 |
+
builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
|
306 |
+
config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small")
|
307 |
+
builder.load_model(selected_model, config)
|
308 |
+
st.session_state['builder'] = builder
|
309 |
st.session_state['model_loaded'] = True
|
310 |
st.rerun()
|
311 |
|
312 |
# Tabs
|
313 |
+
tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs([
|
314 |
+
"Build Titan π±",
|
315 |
+
"Fine-Tune NLP π§ ",
|
316 |
+
"Fine-Tune CV π¨",
|
317 |
+
"Test Titans π§ͺ",
|
318 |
+
"Agentic RAG π",
|
319 |
+
"Camera Inputs π·"
|
320 |
+
])
|
321 |
|
322 |
with tab1:
|
323 |
+
st.header("Build Your Titan π±")
|
324 |
+
model_type = st.selectbox("Model Type", ["NLP (Causal LM)", "CV (Diffusion)"], key="build_type")
|
325 |
+
base_model = st.selectbox(
|
326 |
+
"Select Tiny Model",
|
327 |
+
["HuggingFaceTB/SmolLM-135M", "Qwen/Qwen1.5-0.5B-Chat"] if "NLP" in model_type else ["stabilityai/stable-diffusion-2-1", "CompVis/stable-diffusion-v1-4"]
|
328 |
+
)
|
329 |
model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
|
330 |
if st.button("Download Model β¬οΈ"):
|
331 |
+
config = (ModelConfig if "NLP" in model_type else DiffusionConfig)(name=model_name, base_model=base_model, size="small")
|
332 |
+
builder = ModelBuilder() if "NLP" in model_type else DiffusionBuilder()
|
333 |
builder.load_model(base_model, config)
|
334 |
builder.save_model(config.model_path)
|
335 |
st.session_state['builder'] = builder
|
|
|
337 |
st.rerun()
|
338 |
|
339 |
with tab2:
|
340 |
+
st.header("Fine-Tune NLP Titan π§ (Word Wizardry!)")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
341 |
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], ModelBuilder):
|
342 |
+
st.warning("Load an NLP Titan first! β οΈ")
|
343 |
else:
|
344 |
+
uploaded_csv = st.file_uploader("Upload CSV for NLP SFT", type="csv", key="nlp_csv")
|
345 |
+
if uploaded_csv and st.button("Tune the Wordsmith π§"):
|
346 |
+
csv_path = f"nlp_sft_data_{int(time.time())}.csv"
|
347 |
+
with open(csv_path, "wb") as f:
|
348 |
+
f.write(uploaded_csv.read())
|
349 |
+
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
|
350 |
+
new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
|
351 |
+
st.session_state['builder'].config = new_config
|
352 |
+
st.session_state['builder'].fine_tune_sft(csv_path)
|
353 |
+
st.session_state['builder'].save_model(new_config.model_path)
|
354 |
+
zip_path = f"{new_config.model_path}.zip"
|
355 |
+
zip_directory(new_config.model_path, zip_path)
|
356 |
+
st.markdown(get_download_link(zip_path, "application/zip", "Download Tuned NLP Titan"), unsafe_allow_html=True)
|
357 |
|
358 |
+
with tab3:
|
359 |
+
st.header("Fine-Tune CV Titan π¨ (Vision Vibes!)")
|
360 |
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False) or not isinstance(st.session_state['builder'], DiffusionBuilder):
|
361 |
+
st.warning("Load a CV Titan first! β οΈ")
|
362 |
else:
|
363 |
+
uploaded_files = st.file_uploader("Upload Images/Videos", type=["png", "jpg", "jpeg", "mp4", "mp3"], accept_multiple_files=True, key="cv_upload")
|
364 |
+
text_input = st.text_area("Enter Text (one line per image)", "Batman Neon\nIron Man Hologram\nThor Lightning", key="cv_text")
|
365 |
+
if uploaded_files and st.button("Tune the Visionary ποΈ"):
|
366 |
images = [Image.open(f) for f in uploaded_files if f.type.startswith("image")]
|
367 |
texts = text_input.splitlines()
|
368 |
if len(images) > len(texts):
|
369 |
texts.extend([""] * (len(images) - len(texts)))
|
370 |
elif len(texts) > len(images):
|
371 |
texts = texts[:len(images)]
|
|
|
372 |
st.session_state['builder'].fine_tune_sft(images, texts)
|
373 |
new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
|
374 |
new_config = DiffusionConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small")
|
375 |
st.session_state['builder'].config = new_config
|
376 |
st.session_state['builder'].save_model(new_config.model_path)
|
|
|
377 |
for img, text in zip(images, texts):
|
378 |
filename = generate_filename(text)
|
379 |
img.save(filename)
|
380 |
st.image(img, caption=filename)
|
381 |
zip_path = f"{new_config.model_path}.zip"
|
382 |
zip_directory(new_config.model_path, zip_path)
|
383 |
+
st.markdown(get_download_link(zip_path, "application/zip", "Download Tuned CV Titan"), unsafe_allow_html=True)
|
384 |
+
|
385 |
+
with tab4:
|
386 |
+
st.header("Test Titans π§ͺ (Brains & Eyes!)")
|
387 |
+
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
|
388 |
+
st.warning("Load a Titan first! β οΈ")
|
389 |
+
else:
|
390 |
+
if isinstance(st.session_state['builder'], ModelBuilder):
|
391 |
+
st.subheader("NLP Test π§ ")
|
392 |
+
test_prompt = st.text_area("Enter NLP Prompt", "Plan a superhero party!", key="nlp_test")
|
393 |
+
if st.button("Test NLP Titan βΆοΈ"):
|
394 |
+
result = st.session_state['builder'].evaluate(test_prompt)
|
395 |
+
st.write(f"**Response**: {result}")
|
396 |
+
elif isinstance(st.session_state['builder'], DiffusionBuilder):
|
397 |
+
st.subheader("CV Test π¨")
|
398 |
+
test_prompt = st.text_area("Enter CV Prompt", "Superhero in neon style", key="cv_test")
|
399 |
+
if st.button("Test CV Titan βΆοΈ"):
|
400 |
+
image = st.session_state['builder'].generate(test_prompt)
|
401 |
+
st.image(image, caption="Generated Image")
|
402 |
+
|
403 |
+
cameras = detect_cameras()
|
404 |
+
if cameras:
|
405 |
+
st.subheader("Camera Snapshot Test π·")
|
406 |
+
camera_idx = st.selectbox("Select Camera", cameras, key="camera_select")
|
407 |
+
snapshot_text = st.text_input("Snapshot Text", "Camera Snap", key="snap_text")
|
408 |
+
if st.button("Capture Snapshot πΈ"):
|
409 |
+
cap = cv2.VideoCapture(camera_idx)
|
410 |
+
ret, frame = cap.read()
|
411 |
+
if ret:
|
412 |
+
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
413 |
+
img = Image.fromarray(rgb_frame)
|
414 |
+
filename = generate_filename(snapshot_text)
|
415 |
+
img.save(filename)
|
416 |
+
st.image(img, caption=filename)
|
417 |
+
cap.release()
|
418 |
+
|
419 |
+
with tab5:
|
420 |
+
st.header("Agentic RAG π (Smart Plans & Visions!)")
|
421 |
+
if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
|
422 |
+
st.warning("Load a Titan first! β οΈ")
|
423 |
+
else:
|
424 |
+
if isinstance(st.session_state['builder'], ModelBuilder):
|
425 |
+
st.subheader("NLP RAG Party π§ ")
|
426 |
+
if st.button("Run NLP RAG Demo π"):
|
427 |
+
agent = NLPAgent(st.session_state['builder'].model, st.session_state['builder'].tokenizer)
|
428 |
+
task = "Plan a luxury superhero-themed party at Wayne Manor."
|
429 |
+
plan_df = agent.plan_party(task)
|
430 |
+
st.dataframe(plan_df)
|
431 |
+
elif isinstance(st.session_state['builder'], DiffusionBuilder):
|
432 |
+
st.subheader("CV RAG Enhance π¨")
|
433 |
+
if st.button("Run CV RAG Demo ποΈ"):
|
434 |
+
agent = CVAgent(st.session_state['builder'].pipeline)
|
435 |
+
task = "Enhance superhero images with 2025 trends."
|
436 |
+
enhance_df = agent.enhance_images(task)
|
437 |
+
st.dataframe(enhance_df)
|
438 |
+
|
439 |
+
with tab6:
|
440 |
+
st.header("Camera Inputs π· (Live Feed Fun!)")
|
441 |
+
cameras = detect_cameras()
|
442 |
+
if not cameras:
|
443 |
+
st.warning("No cameras detected! β οΈ")
|
444 |
+
else:
|
445 |
+
st.write(f"Detected {len(cameras)} cameras!")
|
446 |
+
for idx in cameras:
|
447 |
+
st.subheader(f"Camera {idx}")
|
448 |
+
cap = cv2.VideoCapture(idx)
|
449 |
+
if st.button(f"Capture from Camera {idx} πΈ", key=f"cap_{idx}"):
|
450 |
+
ret, frame = cap.read()
|
451 |
+
if ret:
|
452 |
+
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
453 |
+
img = Image.fromarray(rgb_frame)
|
454 |
+
filename = generate_filename(f"Camera_{idx}_snap")
|
455 |
+
img.save(filename)
|
456 |
+
st.image(img, caption=filename)
|
457 |
+
cap.release()
|
458 |
+
|
459 |
+
# Preload demo files
|
460 |
+
demo_images = ["20250319_010000_AM_Batman.png", "20250319_010001_AM_IronMan.png", "20250319_010002_AM_Thor.png"]
|
461 |
+
demo_videos = ["20250319_010000_AM_Batman.mp4", "20250319_010001_AM_IronMan.mp4", "20250319_010002_AM_Thor.mp4"]
|
462 |
+
for img in demo_images:
|
463 |
+
if not os.path.exists(img):
|
464 |
+
Image.new("RGB", (100, 100)).save(img)
|
465 |
+
for vid in demo_videos:
|
466 |
+
if not os.path.exists(vid):
|
467 |
+
with open(vid, "wb") as f:
|
468 |
+
f.write(b"") # Dummy file
|
469 |
+
|
470 |
+
# Demo SFT Dataset
|
471 |
+
st.subheader("Diffusion SFT Demo Dataset π¨")
|
472 |
+
demo_texts = ["Batman Neon", "Iron Man Hologram", "Thor Lightning"]
|
473 |
+
demo_code = "\n".join([f"{i+1}. {text} -> {demo_images[i]}" for i, text in enumerate(demo_texts)])
|
474 |
+
st.code(demo_code, language="text")
|
475 |
+
if st.button("Download Demo CSV π"):
|
476 |
+
csv_path = f"demo_diffusion_sft_{int(time.time())}.csv"
|
477 |
+
with open(csv_path, "w", newline="") as f:
|
478 |
+
writer = csv.writer(f)
|
479 |
+
writer.writerow(["image", "text"])
|
480 |
+
for img, text in zip(demo_images, demo_texts):
|
481 |
+
writer.writerow([img, text])
|
482 |
+
st.markdown(get_download_link(csv_path, "text/csv", "Download Demo CSV"), unsafe_allow_html=True)
|