app.py

import streamlit as st
import os
import py7zr
import requests
from huggingface_hub import HfApi
import torch
from torch.utils.data import DataLoader
import shutil
from pathlib import Path
from typing import Optional
import sys
import io

# Import the denoising code (assuming it's in a file called denoising_model.py)
from denoising_model import DenoisingModel, DenoiseDataset, get_optimal_threads

class StreamCapture:
    def __init__(self):
        self.logs = []
    
    def write(self, text):
        self.logs.append(text)
        st.warning(text)
    
    def flush(self):
        pass

def download_and_extract_7z(url: str, extract_to: str = '.') -> Optional[str]:
    """Downloads a 7z file and extracts it"""
    try:
        st.warning(f"Downloading file from {url}...")
        response = requests.get(url, stream=True)
        response.raise_for_status()
        
        archive_path = os.path.join(extract_to, 'dataset.7z')
        with open(archive_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        st.warning("Extracting 7z archive...")
        with py7zr.SevenZipFile(archive_path, mode='r') as z:
            z.extractall(extract_to)
        
        # Handle directory renaming
        output_images_path = os.path.join(extract_to, 'output_images')
        if os.path.exists(output_images_path):
            # Move and rename directories
            source_noisy = os.path.join(output_images_path, 'images_noisy')
            source_target = os.path.join(output_images_path, 'images_target')
            
            if os.path.exists('noisy_images'):
                shutil.rmtree('noisy_images')
            if os.path.exists('target_images'):
                shutil.rmtree('target_images')
                
            shutil.move(source_noisy, 'noisy_images')
            shutil.move(source_target, 'target_images')
            
            # Clean up
            if os.path.exists(output_images_path):
                shutil.rmtree(output_images_path)
            
        os.remove(archive_path)
        st.warning("Download and extraction completed successfully.")
        return None
        
    except Exception as e:
        return f"Error: {str(e)}"

def upload_to_huggingface(file_path: str, repo_id: str, path_in_repo: str):
    """Uploads a file to Hugging Face"""
    try:
        api = HfApi()
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=path_in_repo,
            repo_id=repo_id,
            repo_type="space"
        )
        st.warning(f"Successfully uploaded {file_path} to {repo_id}")
    except Exception as e:
        st.error(f"Error uploading to Hugging Face: {str(e)}")

def train_model_with_upload(epochs, batch_size, learning_rate, save_interval, num_workers, repo_id):
    """Modified training function that uploads checkpoints to Hugging Face"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    st.warning(f"Using device: {device}")
    
    # Create temporary directory for checkpoints
    checkpoint_dir = "temp_checkpoints"
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    try:
        dataset = DenoiseDataset('noisy_images', 'target_images')
        dataloader = DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=num_workers,
            pin_memory=True if torch.cuda.is_available() else False
        )
        
        model = DenoisingModel().to(device)
        criterion = torch.nn.MSELoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
        
        for epoch in range(epochs):
            st.warning(f"Starting epoch {epoch+1}/{epochs}")
            for batch_idx, (noisy_patches, target_patches) in enumerate(dataloader):
                noisy_patches = noisy_patches.to(device)
                target_patches = target_patches.to(device)
                
                outputs = model(noisy_patches)
                loss = criterion(outputs, target_patches)
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                if (batch_idx + 1) % 10 == 0:
                    st.warning(f"Epoch [{epoch+1}/{epochs}], Batch [{batch_idx+1}], Loss: {loss.item():.6f}")
                
                if (batch_idx + 1) % save_interval == 0:
                    checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint_epoch{epoch+1}_batch{batch_idx+1}.pth")
                    torch.save(model.state_dict(), checkpoint_path)
                    
                    # Upload checkpoint to Hugging Face
                    upload_to_huggingface(
                        checkpoint_path,
                        repo_id,
                        f"checkpoints/checkpoint_epoch{epoch+1}_batch{batch_idx+1}.pth"
                    )
        
        # Save and upload final model
        final_model_path = os.path.join(checkpoint_dir, "final_model.pth")
        torch.save(model.state_dict(), final_model_path)
        upload_to_huggingface(final_model_path, repo_id, "model/final_model.pth")
        
    finally:
        # Clean up temporary directory
        if os.path.exists(checkpoint_dir):
            shutil.rmtree(checkpoint_dir)

def main():
    st.title("Image Denoising Model Training")
    
    # Redirect stdout to capture print statements
    sys.stdout = StreamCapture()
    
    # Input for Hugging Face token
    hf_token = st.text_input("Enter your Hugging Face token:", type="password")
    if hf_token:
        os.environ["HF_TOKEN"] = hf_token
    
    # Input for repository ID
    repo_id = st.text_input("Enter your Hugging Face repository ID (username/repo):")
    
    # Download and extract dataset button
    if st.button("Download and Extract Dataset"):
        url = "https://huggingface.co/spaces/vericudebuget/ok4231/resolve/main/output_images.7z"
        error = download_and_extract_7z(url)
        if error:
            st.error(error)
    
    # Training parameters
    col1, col2 = st.columns(2)
    with col1:
        epochs = st.number_input("Number of epochs", min_value=1, value=10)
        batch_size = st.number_input("Batch size", min_value=1, value=4)
        learning_rate = st.number_input("Learning rate", min_value=0.0001, value=0.001, format="%.4f")
    
    with col2:
        save_interval = st.number_input("Save interval (batches)", min_value=1, value=1000)
        num_workers = st.number_input("Number of workers", min_value=1, value=get_optimal_threads())
    
    # Start training button
    if st.button("Start Training"):
        if not hf_token:
            st.error("Please enter your Hugging Face token")
            return
        if not repo_id:
            st.error("Please enter your repository ID")
            return
        if not os.path.exists("noisy_images") or not os.path.exists("target_images"):
            st.error("Dataset not found. Please download and extract it first.")
            return
            
        train_model_with_upload(epochs, batch_size, learning_rate, save_interval, num_workers, repo_id)

if __name__ == "__main__":
    main()