Update app.py

app.py

@@ -10,7 +10,6 @@ import os
 import time
 import threading
 from PIL import Image
-import io
 import numpy as np
 
 # ======================
@@ -23,7 +22,7 @@ CONFIG = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
         },
         "max_images": 100,
-        "scrape_time": 10  # 3 hours in seconds (simulated for testing)
+        "scrape_time": 10
     },
     "training": {
         "batch_size": 4,
@@ -46,39 +45,48 @@ class WebScraper:
     def __init__(self):
         self.stop_event = threading.Event()
         self.scraped_data = []
+        self._lock = threading.Lock()
         
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state['stop_event']
+        del state['_lock']
+        return state
+        
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.stop_event = threading.Event()
+        self._lock = threading.Lock()
+
     def scrape_images(self, query):
-        search_url = CONFIG["scraping"]["search_url"].format(query=query)
-        try:
-            response = requests.get(search_url, headers=CONFIG["scraping"]["headers"])
-            soup = BeautifulSoup(response.content, 'html.parser')
-            
-            # Extract image URLs (example selector - needs adjustment for actual site)
-            img_tags = soup.find_all('img', {'class': 'photo-item__img'})
-            for img in img_tags[:CONFIG["scraping"]["max_images"]]:
-                if self.stop_event.is_set():
-                    break
-                img_url = img['src']
-                try:
-                    img_data = requests.get(img_url).content
-                    img_name = f"{int(time.time())}.jpg"
-                    img_path = os.path.join(CONFIG["paths"]["dataset_dir"], img_name)
-                    
-                    with open(img_path, 'wb') as f:
-                        f.write(img_data)
-                    
-                    # Store text-image pair (text = query)
-                    self.scraped_data.append({"text": query, "image": img_path})
-                except Exception as e:
-                    print(f"Error downloading image: {e}")
-        except Exception as e:
-            print(f"Scraping error: {e}")
+        with self._lock:
+            search_url = CONFIG["scraping"]["search_url"].format(query=query)
+            try:
+                response = requests.get(search_url, headers=CONFIG["scraping"]["headers"])
+                soup = BeautifulSoup(response.content, 'html.parser')
+                img_tags = soup.find_all('img', {'class': 'photo-item__img'})
+                
+                for img in img_tags[:CONFIG["scraping"]["max_images"]]:
+                    if self.stop_event.is_set():
+                        break
+                    img_url = img['src']
+                    try:
+                        img_data = requests.get(img_url).content
+                        img_name = f"{int(time.time())}.jpg"
+                        img_path = os.path.join(CONFIG["paths"]["dataset_dir"], img_name)
+                        
+                        with open(img_path, 'wb') as f:
+                            f.write(img_data)
+                            
+                        self.scraped_data.append({"text": query, "image": img_path})
+                    except Exception as e:
+                        print(f"Error downloading image: {e}")
+            except Exception as e:
+                print(f"Scraping error: {e}")
             
     def start_scraping(self, query):
         self.stop_event.clear()
-        if not os.path.exists(CONFIG["paths"]["dataset_dir"]):
-            os.makedirs(CONFIG["paths"]["dataset_dir"])
-        
+        os.makedirs(CONFIG["paths"]["dataset_dir"], exist_ok=True)
         thread = threading.Thread(target=self.scrape_images, args=(query,))
         thread.start()
         return "Scraping started..."
@@ -97,17 +105,12 @@ class TextImageDataset(Dataset):
     def __getitem__(self, idx):
         item = self.data[idx]
         image = Image.open(item["image"]).convert('RGB')
-        
-        if self.transform:
-            image = self.transform(image)
-            
         return {"text": item["text"], "image": image}
 
-# Simplified Text-to-Image Generator
 class TextConditionedGenerator(nn.Module):
     def __init__(self):
         super().__init__()
-        self.text_embedding = nn.Embedding(1000, 128)  # Simplified text embedding
+        self.text_embedding = nn.Embedding(1000, 128)
         self.model = nn.Sequential(
             nn.Linear(128 + CONFIG["training"]["latent_dim"], 256),
             nn.LeakyReLU(0.2),
@@ -121,8 +124,7 @@ class TextConditionedGenerator(nn.Module):
     def forward(self, text, noise):
         text_emb = self.text_embedding(text)
         combined = torch.cat([text_emb, noise], 1)
-        img = self.model(combined)
-        return img.view(-1, 3, CONFIG["training"]["img_size"], CONFIG["training"]["img_size"])
+        return self.model(combined).view(-1, 3, CONFIG["training"]["img_size"], CONFIG["training"]["img_size"])
 
 # ======================
 # Training Utilities
@@ -144,29 +146,26 @@ def train_model(scraper, progress=gr.Progress()):
     criterion = nn.BCELoss()
     
     for epoch in progress.tqdm(range(CONFIG["training"]["epochs"]), desc="Training"):
-        for i, batch in enumerate(dataloader):
-            # Train discriminator
+        for batch in dataloader:
             real_imgs = batch["image"]
             real_labels = torch.ones(real_imgs.size(0), 1)
-            
             noise = torch.randn(real_imgs.size(0), CONFIG["training"]["latent_dim"])
-            fake_imgs = generator(torch.randint(0, 1000, (real_imgs.size(0),)), noise)
-            fake_labels = torch.zeros(real_imgs.size(0), 1)
             
+            # Discriminator training
             optimizer_D.zero_grad()
             real_loss = criterion(discriminator(real_imgs.view(-1, 3*64**2)), real_labels)
-            fake_loss = criterion(discriminator(fake_imgs.detach().view(-1, 3*64**2)), fake_labels)
-            d_loss = real_loss + fake_loss
+            fake_imgs = generator(torch.randint(0, 1000, (real_imgs.size(0),)), noise)
+            fake_loss = criterion(discriminator(fake_imgs.detach().view(-1, 3*64**2)), torch.zeros_like(real_labels))
+            d_loss = (real_loss + fake_loss) / 2
             d_loss.backward()
             optimizer_D.step()
             
-            # Train generator
+            # Generator training
             optimizer_G.zero_grad()
-            validity = discriminator(fake_imgs.view(-1, 3*64**2))
-            g_loss = criterion(validity, torch.ones_like(validity))
+            g_loss = criterion(discriminator(fake_imgs.view(-1, 3*64**2)), torch.ones_like(real_labels))
             g_loss.backward()
             optimizer_G.step()
-            
+    
     torch.save(generator.state_dict(), CONFIG["paths"]["model_save"])
     return "Training completed!"
 
@@ -179,16 +178,14 @@ class ModelRunner:
         self.custom_model = None
         
     def load_pretrained(self):
-        if self.pretrained_pipe is None:
-            self.pretrained_pipe = DiffusionPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-xl-base-1.0"
-            )
+        if not self.pretrained_pipe:
+            self.pretrained_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
         return self.pretrained_pipe
     
     def load_custom(self):
-        if self.custom_model is None:
+        if not self.custom_model:
             model = TextConditionedGenerator()
-            model.load_state_dict(torch.load(CONFIG["paths"]["model_save"]))
+            model.load_state_dict(torch.load(CONFIG["paths"]["model_save"], map_location='cpu'))
             self.custom_model = model
         return self.custom_model
 
@@ -196,42 +193,39 @@ class ModelRunner:
 # Gradio Interface
 # ======================
 with gr.Blocks() as app:
-    # Use Gradio's state management
-    scraper_state = gr.State(WebScraper())
-    model_runner_state = gr.State(ModelRunner())
+    scraper_state = gr.State(WebScraper)
+    model_runner_state = gr.State(ModelRunner)
     
     with gr.Row():
         with gr.Column():
             query_input = gr.Textbox(label="Search Query")
             scrape_btn = gr.Button("Start Scraping")
             scrape_status = gr.Textbox(label="Scraping Status")
-            
             train_btn = gr.Button("Start Training")
             training_status = gr.Textbox(label="Training Status")
             
         with gr.Column():
             prompt_input = gr.Textbox(label="Generation Prompt")
-            model_choice = gr.Radio(["Pretrained", "Custom"], label="Model Type")
+            model_choice = gr.Radio(["Pretrained", "Custom"], label="Model Type", value="Pretrained")
             generate_btn = gr.Button("Generate Image")
             output_image = gr.Image(label="Generated Image")
-    
-    # Event Handlers
+
     scrape_btn.click(
-        fn=lambda scraper, query: scraper.start_scraping(query),
-        inputs=[scraper_state, query_input],
-        outputs=scrape_status
+        lambda scraper, query: scraper.start_scraping(query),
+        [scraper_state, query_input],
+        scrape_status
     )
     
     train_btn.click(
-        fn=lambda scraper: train_model(scraper),
-        inputs=[scraper_state],
-        outputs=training_status
+        lambda scraper: train_model(scraper),
+        [scraper_state],
+        training_status
     )
     
     generate_btn.click(
-        fn=lambda prompt, model_type, runner: generate_image(prompt, model_type, runner),
-        inputs=[prompt_input, model_choice, model_runner_state],
-        outputs=output_image
+        lambda prompt, model_type, runner: generate_image(prompt, model_type, runner),
+        [prompt_input, model_choice, model_runner_state],
+        output_image
     )
 
 def generate_image(prompt, model_type, runner):
@@ -240,12 +234,11 @@ def generate_image(prompt, model_type, runner):
         image = pipe(prompt).images[0]
     else:
         model = runner.load_custom()
-        # Simplified generation process
         noise = torch.randn(1, CONFIG["training"]["latent_dim"])
-        fake = model(torch.randint(0, 1000, (1,)), noise).detach()
-        image = fake.squeeze().permute(1,2,0).numpy()
-        image = (image + 1) / 2  # Scale to [0,1]
-    
+        with torch.no_grad():
+            fake = model(torch.randint(0, 1000, (1,)), noise)
+        image = fake.squeeze().permute(1, 2, 0).numpy()
+        image = (image + 1) / 2
     return Image.fromarray((image * 255).astype(np.uint8))
 
 if __name__ == "__main__":