Update app.py

app.py

@@ -19,7 +19,7 @@ CONFIG = {
     "scraping": {
         "search_url": "https://www.pexels.com/search/{query}/",
         "headers": {
-            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
         },
         "max_images": 100,
         "progress_interval": 1
@@ -77,7 +77,10 @@ class WebScraper:
                     if self.stop_event.is_set():
                         break
 
-                    img_url = img['src']
+                    img_url = img.get('src')
+                    if not img_url:
+                        continue
+
                     try:
                         img_data = requests.get(img_url).content
                         img_name = f"{int(time.time())}_{idx}.jpg"
@@ -97,6 +100,7 @@ class WebScraper:
                 self.scraping_progress = 100
 
     def start_scraping(self, query):
+        self.scraped_data.clear()
         self.stop_event.clear()
         thread = threading.Thread(target=self.scrape_images, args=(query,))
         thread.start()
@@ -114,8 +118,14 @@ class TextImageDataset(Dataset):
 
     def __getitem__(self, idx):
         item = self.data[idx]
-        image = Image.open(item["image"]).convert('RGB')
-        image = torch.randn(3, 64, 64)  # Simplified for example
+        try:
+            image = Image.open(item["image"]).convert('RGB')
+            image = image.resize((64, 64))
+            image = np.array(image).transpose(2, 0, 1) / 127.5 - 1
+            image = torch.tensor(image, dtype=torch.float32)
+        except Exception as e:
+            print(f"Error loading image: {e}")
+            image = torch.randn(3, 64, 64)
         return {"text": item["text"], "image": image}
 
 class TextConditionedGenerator(nn.Module):
@@ -159,23 +169,26 @@ def train_model(scraper, progress=gr.Progress()):
     optimizer_D = optim.Adam(discriminator.parameters(), lr=CONFIG["training"]["lr"])
     criterion = nn.BCELoss()
 
-    total_batches = len(dataloader)
     for epoch in progress.tqdm(range(CONFIG["training"]["epochs"])):
-        for batch_idx, batch in enumerate(dataloader):
-            real_imgs = torch.randn(4, 3, 64, 64)
-            real_labels = torch.ones(real_imgs.size(0), 1)
+        for batch in dataloader:
+            real_imgs = batch["image"]
+            text_tokens = torch.randint(0, 1000, (real_imgs.size(0),))
             noise = torch.randn(real_imgs.size(0), 100)
+            real_labels = torch.ones(real_imgs.size(0), 1)
+            fake_labels = torch.zeros(real_imgs.size(0), 1)
 
+            # Discriminator update
             optimizer_D.zero_grad()
             real_loss = criterion(discriminator(real_imgs.view(-1, 3*64*64)), real_labels)
-            fake_imgs = generator(torch.randint(0, 1000, (real_imgs.size(0),)), noise)
-            fake_loss = criterion(discriminator(fake_imgs.detach().view(-1, 3*64*64)), torch.zeros_like(real_labels))
+            fake_imgs = generator(text_tokens, noise)
+            fake_loss = criterion(discriminator(fake_imgs.detach().view(-1, 3*64*64)), fake_labels)
             d_loss = (real_loss + fake_loss) / 2
             d_loss.backward()
             optimizer_D.step()
 
+            # Generator update
             optimizer_G.zero_grad()
-            g_loss = criterion(discriminator(fake_imgs.view(-1, 3*64*64)), torch.ones_like(real_labels))
+            g_loss = criterion(discriminator(fake_imgs.view(-1, 3*64*64)), real_labels)
             g_loss.backward()
             optimizer_G.step()