awacke1 commited on
Commit
3820640
Β·
verified Β·
1 Parent(s): 1dee6c7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +242 -193
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 # For diffusion models
23
  from urllib.parse import quote
 
24
 
25
- # Set up logging
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 Configuration Classes
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 Builder Classes
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
- with st.spinner(f"Loading {model_path}... ⏳"):
111
- self.model = AutoModelForCausalLM.from_pretrained(model_path)
112
- self.tokenizer = AutoTokenizer.from_pretrained(model_path)
113
- if self.tokenizer.pad_token is None:
114
- self.tokenizer.pad_token = self.tokenizer.eos_token
115
- if config:
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
- with st.spinner(f"Training epoch {epoch + 1}/{epochs}... βš™οΈ"):
136
- total_loss = 0
137
- for batch in dataloader:
138
- optimizer.zero_grad()
139
- input_ids = batch["input_ids"].to(device)
140
- attention_mask = batch["attention_mask"].to(device)
141
- labels = batch["labels"].to(device)
142
- outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
143
- loss = outputs.loss
144
- loss.backward()
145
- optimizer.step()
146
- total_loss += loss.item()
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
- with st.spinner("Saving model... πŸ’Ύ"):
153
- os.makedirs(os.path.dirname(path), exist_ok=True)
154
- self.model.save_pretrained(path)
155
- self.tokenizer.save_pretrained(path)
156
- st.success(f"Model saved at {path}! βœ…")
157
-
158
- def evaluate(self, prompt: str, status_container=None):
159
  self.model.eval()
160
- if status_container:
161
- status_container.write("Preparing to evaluate... 🧠")
162
- try:
163
- with torch.no_grad():
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
- with st.spinner(f"Loading diffusion model {model_path}... ⏳"):
179
- self.pipeline = StableDiffusionPipeline.from_pretrained(model_path)
180
- self.pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
181
- if config:
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
- with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... βš™οΈ"):
194
- total_loss = 0
195
- for batch in dataloader:
196
- optimizer.zero_grad()
197
- image = batch["image"].to(self.pipeline.device)
198
- text = batch["text"]
199
- latents = self.pipeline.vae.encode(image).latent_dist.sample()
200
- noise = torch.randn_like(latents)
201
- timesteps = torch.randint(0, self.pipeline.scheduler.num_train_timesteps, (latents.shape[0],), device=latents.device)
202
- noisy_latents = self.pipeline.scheduler.add_noise(latents, noise, timesteps)
203
- text_embeddings = self.pipeline.text_encoder(self.pipeline.tokenizer(text, return_tensors="pt").input_ids.to(self.pipeline.device))[0]
204
- pred_noise = self.pipeline.unet(noisy_latents, timesteps, encoder_hidden_states=text_embeddings).sample
205
- loss = torch.nn.functional.mse_loss(pred_noise, noise)
206
- loss.backward()
207
- optimizer.step()
208
- total_loss += loss.item()
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
- with st.spinner("Saving diffusion model... πŸ’Ύ"):
215
- os.makedirs(os.path.dirname(path), exist_ok=True)
216
- self.pipeline.save_pretrained(path)
217
- st.success(f"Diffusion model saved at {path}! βœ…")
218
 
219
- # Utility Functions
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
- file_path = os.path.join(root, file)
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
- files = []
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
- links_md = ' '.join([f"[{name}]({url})" for name, url in search_urls.items()])
258
- return links_md
259
-
260
- # Agent Class
261
- class PartyPlannerAgent:
 
 
 
 
 
 
 
 
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
- locations = {"Wayne Manor": (42.3601, -71.0589), "New York": (40.7128, -74.0060), "Los Angeles": (34.0522, -118.2437), "London": (51.5074, -0.1278)}
282
- wayne_coords = locations["Wayne Manor"]
283
- travel_times = {loc: calculate_cargo_travel_time(coords, wayne_coords) for loc, coords in locations.items() if loc != "Wayne Manor"}
284
-
 
 
 
 
 
 
 
 
 
 
 
 
285
  data = [
286
- {"Location": "New York", "Travel Time (hrs)": travel_times["New York"], "Luxury Idea": "Gold-plated Batman statues"},
287
- {"Location": "Los Angeles", "Travel Time (hrs)": travel_times["Los Angeles"], "Luxury Idea": "VR superhero battles"},
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
- ("Image Gallery πŸ“Έ", ["png", "jpg", "jpeg"]),
315
- ("Video Gallery πŸŽ₯", ["mp4"]),
316
- ("Audio Gallery 🎢", ["mp3"])
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 "Image" in gallery_type:
326
  st.image(Image.open(file), caption=file, use_column_width=True)
327
- elif "Video" in gallery_type:
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 model_type == "Causal LM" else "diffusion")
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 'builder' not in st.session_state:
338
- st.session_state['builder'] = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
339
- config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=os.path.basename(selected_model), base_model="unknown", size="small")
340
- st.session_state['builder'].load_model(selected_model, config)
341
  st.session_state['model_loaded'] = True
342
  st.rerun()
343
 
344
  # Tabs
345
- tab1, tab2, tab3, tab4, tab5 = st.tabs(["Build Tiny Titan 🌱", "Fine-Tune Titan πŸ”§", "Test Titan πŸ§ͺ", "Agentic RAG Party 🌐", "Diffusion SFT 🎨"])
 
 
 
 
 
 
 
346
 
347
  with tab1:
348
- st.header("Build Tiny Titan 🌱")
349
- model_type = st.selectbox("Model Type", ["Causal LM", "Diffusion"], key="build_type")
350
- if model_type == "Causal LM":
351
- base_model = st.selectbox("Select Tiny Model", ["HuggingFaceTB/SmolLM-135M", "HuggingFaceTB/SmolLM-360M", "Qwen/Qwen1.5-0.5B-Chat"])
352
- else:
353
- base_model = st.selectbox("Select Tiny Diffusion Model", ["stabilityai/stable-diffusion-2-1", "runwayml/stable-diffusion-v1-5", "CompVis/stable-diffusion-v1-4"])
354
  model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
355
  if st.button("Download Model ⬇️"):
356
- config = (ModelConfig if model_type == "Causal LM" else DiffusionConfig)(name=model_name, base_model=base_model, size="small")
357
- builder = ModelBuilder() if model_type == "Causal LM" else DiffusionBuilder()
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("Please build or load a Causal LM Titan first! ⚠️")
399
  else:
400
- if st.button("Run Agentic RAG Demo πŸŽ‰"):
401
- agent = PartyPlannerAgent(model=st.session_state['builder'].model, tokenizer=st.session_state['builder'].tokenizer)
402
- task = "Plan a luxury superhero-themed party at Wayne Manor."
403
- plan_df = agent.plan_party(task)
404
- st.dataframe(plan_df)
 
 
 
 
 
 
 
 
405
 
406
- with tab5:
407
- st.header("Diffusion SFT 🎨")
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("Please build or load a Diffusion Titan first! ⚠️")
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)", "Line 1\nLine 2\nLine 3")
413
- if uploaded_files and st.button("Fine-Tune Diffusion Model πŸ”„"):
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 Fine-Tuned Diffusion Model"), unsafe_allow_html=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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)