app.py

import gradio as gr
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from diffusers import DiffusionPipeline
import requests
from bs4 import BeautifulSoup
import os
import time
import threading
from PIL import Image
import numpy as np

# ======================
# Configuration
# ======================
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"
        },
        "max_images": 100,
        "progress_interval": 1
    },
    "training": {
        "batch_size": 4,
        "epochs": 10,
        "lr": 0.0002,
        "latent_dim": 100,
        "img_size": 64,
        "num_workers": 0,
        "progress_interval": 0.5
    },
    "paths": {
        "dataset_dir": "scraped_data",
        "model_save": "text2img_model.pth"
    }
}

# ======================
# Web Scraping Module
# ======================
class WebScraper:
    def __init__(self):
        self.stop_event = threading.Event()
        self.scraped_data = []
        self._lock = threading.Lock()
        self.scraping_progress = 0
        self.scraped_count = 0
        self.total_images = 0
        
    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):
        with self._lock:
            self.scraping_progress = 0
            self.scraped_count = 0
            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'})
                self.total_images = min(len(img_tags), CONFIG["scraping"]["max_images"])
                
                for idx, img in enumerate(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())}_{idx}.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})
                        self.scraped_count = idx + 1
                        self.scraping_progress = (idx + 1) / self.total_images * 100
                        
                    except Exception as e:
                        print(f"Error downloading image: {e}")
                    
                    time.sleep(0.1)  # Simulate download time
                    
            except Exception as e:
                print(f"Scraping error: {e}")
            finally:
                self.scraping_progress = 100
                
    def start_scraping(self, query):
        self.stop_event.clear()
        os.makedirs(CONFIG["paths"]["dataset_dir"], exist_ok=True)
        thread = threading.Thread(target=self.scrape_images, args=(query,))
        thread.start()
        return "Scraping started..."

# ======================
# Dataset and Models
# ======================
class TextImageDataset(Dataset):
    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        item = self.data[idx]
        image = Image.open(item["image"]).convert('RGB')
        image = torch.randn(3, 64, 64)  # Simplified for example
        return {"text": item["text"], "image": image}

class TextConditionedGenerator(nn.Module):
    def __init__(self):
        super().__init__()
        self.text_embedding = nn.Embedding(1000, 128)
        self.model = nn.Sequential(
            nn.Linear(128 + 100, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 3*64*64),
            nn.Tanh()
        )
        
    def forward(self, text, noise):
        text_emb = self.text_embedding(text)
        combined = torch.cat([text_emb, noise], 1)
        return self.model(combined).view(-1, 3, 64, 64)

# ======================
# Training Utilities
# ======================
def train_model(scraper, progress=gr.Progress()):
    if len(scraper.scraped_data) == 0:
        return "Error: No images scraped! Scrape images first."
    
    dataset = TextImageDataset(scraper.scraped_data)
    dataloader = DataLoader(dataset, 
                          batch_size=CONFIG["training"]["batch_size"], 
                          shuffle=True)
    
    generator = TextConditionedGenerator()
    discriminator = nn.Sequential(
        nn.Linear(3*64*64, 512),
        nn.LeakyReLU(0.2),
        nn.Linear(512, 1),
        nn.Sigmoid()
    )
    
    optimizer_G = optim.Adam(generator.parameters(), lr=CONFIG["training"]["lr"])
    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"]), desc="Epochs"):
        for batch_idx, batch in enumerate(dataloader):
            real_imgs = torch.randn(4, 3, 64, 64)  # Simplified data
            real_labels = torch.ones(real_imgs.size(0), 1)
            noise = torch.randn(real_imgs.size(0), 100)
            
            # Train discriminator
            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))
            d_loss = (real_loss + fake_loss) / 2
            d_loss.backward()
            optimizer_D.step()
            
            # Train generator
            optimizer_G.zero_grad()
            g_loss = criterion(discriminator(fake_imgs.view(-1, 3*64*64)), torch.ones_like(real_labels))
            g_loss.backward()
            optimizer_G.step()
            
            progress(
                (epoch + (batch_idx+1)/total_batches) / CONFIG["training"]["epochs"],
                desc=f"Epoch {epoch+1} | Batch {batch_idx+1}/{total_batches}",
                unit="epoch"
            )
    
    torch.save(generator.state_dict(), CONFIG["paths"]["model_save"])
    return f"Training complete! Used {len(dataset)} samples"

# ======================
# Gradio Interface
# ======================
def create_interface():
    with gr.Blocks() as app:
        scraper = gr.State(WebScraper)
        model_runner = 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")
                scraping_progress = gr.Textbox(label="Scraping Progress", value="0% (0/0)")
                
                train_btn = gr.Button("Start Training")
                training_status = gr.Textbox(label="Training Status")
                training_progress = gr.Textbox(label="Training Progress", value="Not started")

            with gr.Column():
                prompt_input = gr.Textbox(label="Generation Prompt")
                model_choice = gr.Radio(["Pretrained", "Custom"], label="Model Type", value="Pretrained")
                generate_btn = gr.Button("Generate Image")
                output_image = gr.Image(label="Generated Image")

        # Real-time updates using event triggers
        def update_scraping_progress(scraper):
            return f"{scraper.scraping_progress:.1f}% ({scraper.scraped_count}/{scraper.total_images})"
        
        def update_training_progress():
            if os.path.exists(CONFIG["paths"]["model_save"]):
                stats = os.stat(CONFIG["paths"]["model_save"])
                return f"Model size: {stats.st_size//1024}KB"
            return "No trained model"

        # Set up periodic updates
        scraping_progress.change(
            update_scraping_progress,
            inputs=[scraper],
            outputs=[scraping_progress],
            every=CONFIG["scraping"]["progress_interval"]
        )
        
        training_progress.change(
            update_training_progress,
            outputs=[training_progress],
            every=CONFIG["training"]["progress_interval"]
        )

        # Event handlers
        scrape_btn.click(
            lambda s, q: s.start_scraping(q),
            [scraper, query_input],
            [scrape_status]
        )
        
        train_btn.click(
            lambda s: train_model(s),
            [scraper],
            [training_status, training_progress]
        )
        
        generate_btn.click(
            lambda p, m, r: generate_image(p, m, r),
            [prompt_input, model_choice, model_runner],
            [output_image]
        )
        
    return app

def generate_image(prompt, model_type, runner):
    if model_type == "Pretrained":
        pipe = runner.load_pretrained()
        image = pipe(prompt).images[0]
    else:
        model = runner.load_custom()
        noise = torch.randn(1, 100)
        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))

class ModelRunner:
    def __init__(self):
        self.pretrained_pipe = None
        
    def load_pretrained(self):
        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):
        model = TextConditionedGenerator()
        model.load_state_dict(torch.load(CONFIG["paths"]["model_save"], map_location='cpu'))
        return model

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()