Update app_bck.py

app_bck.py

@@ -0,0 +1,399 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+from huggingface_hub import Repository
+from huggingface_hub import HfApi, HfFolder, Repository, create_repo
+import os
+import pandas as pd
+import gradio as gr
+from PIL import Image
+import numpy as np
+from small_256_model import UNet as small_UNet
+from big_1024_model import UNet as big_UNet
+from CLIP import load as load_clip
+from rich import print as rp
+
+# Device configuration
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+big = False if device == torch.device('cpu') else True
+
+# Parameters
+IMG_SIZE = 1024 if big else 256
+BATCH_SIZE = 1 if big else 1
+EPOCHS = 12
+LR = 0.0002
+dataset_id = "K00B404/pix2pix_flux_set"
+model_repo_id = "K00B404/pix2pix_flux"
+
+# Global model variable
+global_model = None
+
+# CLIP
+clip_model, clip_tokenizer = load_clip()
+
+def load_model():
+    """Load the models at startup"""
+    global global_model
+    weights_name = 'big_model_weights.pth' if big else 'small_model_weights.pth'
+    try:
+        checkpoint = torch.load(weights_name, map_location=device)
+        model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.to(device)
+        model.eval()
+        global_model = model
+        rp("Model loaded successfully!")
+        return model
+    except Exception as e:
+        rp(f"Error loading model: {e}")
+        model = big_UNet().to(device) if big else small_UNet().to(device)
+        global_model = model
+        return model
+
+class Pix2PixDataset(torch.utils.data.Dataset):
+    def __init__(self, combined_data, transform, clip_tokenizer):
+        self.data = combined_data
+        self.transform = transform
+        self.clip_tokenizer = clip_tokenizer
+        self.original_folder = 'images_dataset/original/'
+        self.target_folder = 'images_dataset/target/'
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        original_img_filename = os.path.basename(self.data.iloc[idx]['image_path'])
+        original_img_path = os.path.join(self.original_folder, original_img_filename)
+        target_img_path = os.path.join(self.target_folder, original_img_filename)
+
+        original_img = Image.open(original_img_path).convert('RGB')
+        target_img = Image.open(target_img_path).convert('RGB')
+        
+        # Transform images
+        original = self.transform(original_img)
+        target = self.transform(target_img)
+        
+        # Get prompts from the DataFrame
+        original_prompt = self.data.iloc[idx]['original_prompt']
+        enhanced_prompt = self.data.iloc[idx]['enhanced_prompt']
+        
+        # Tokenize the prompts using CLIP tokenizer
+        original_tokens = self.clip_tokenizer(original_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
+        enhanced_tokens = self.clip_tokenizer(enhanced_prompt, return_tensors="pt", padding=True, truncation=True, max_length=77)
+        
+        return original, target, original_tokens, enhanced_tokens
+
+
+
+class UNetWrapper:
+    def __init__(self, unet_model, repo_id, epoch, loss, optimizer, scheduler=None):
+        self.loss = loss
+        self.epoch = epoch
+        self.model = unet_model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.repo_id = repo_id
+        self.token = os.getenv('NEW_TOKEN')  # Ensure the token is set in the environment
+        self.api = HfApi(token=self.token)
+
+    def save_checkpoint(self, save_path):
+        """Save checkpoint with model, optimizer, and scheduler states."""
+        self.save_dict = {
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
+            'model_config': {
+                'big': isinstance(self.model, big_UNet),
+                'img_size': 1024 if isinstance(self.model, big_UNet) else 256
+            },
+            'epoch': self.epoch,
+            'loss': self.loss
+        }
+        torch.save(self.save_dict, save_path)
+        print(f"Checkpoint saved at epoch {self.epoch}, loss: {self.loss}")
+
+    def load_checkpoint(self, checkpoint_path):
+        """Load model, optimizer, and scheduler states from the checkpoint."""
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        if self.scheduler and checkpoint['scheduler_state_dict']:
+            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        self.epoch = checkpoint['epoch']
+        self.loss = checkpoint['loss']
+        print(f"Checkpoint loaded: epoch {self.epoch}, loss: {self.loss}")
+
+    def push_to_hub(self, pth_name):
+        """Push model checkpoint and metadata to the Hugging Face Hub."""
+        try:
+            self.api.upload_file(
+                path_or_fileobj=pth_name,
+                path_in_repo=pth_name,
+                repo_id=self.repo_id,
+                token=self.token,
+                repo_type="model"
+            )
+            print(f"Model checkpoint successfully uploaded to {self.repo_id}")
+        except Exception as e:
+            print(f"Error uploading model: {e}")
+
+        
+   
+            
+        # Create and upload model card
+        model_card = f"""---
+tags:
+- unet
+- pix2pix
+- pytorch
+library_name: pytorch
+license: wtfpl
+datasets:
+- K00B404/pix2pix_flux_set
+language:
+- en
+pipeline_tag: image-to-image
+---
+# Pix2Pix UNet Model
+## Model Description
+Custom UNet model for Pix2Pix image translation.
+- **Image Size:** {self.save_dict['model_config']['img_size']}
+- **Model Type:** {"big" if big else "small"}_UNet ({self.save_dict['model_config']['img_size']})
+## Usage
+```python
+import torch
+from small_256_model import UNet as small_UNet
+from big_1024_model import UNet as big_UNet
+big = True
+# Load the model
+name='big_model_weights.pth' if big else 'small_model_weights.pth'
+checkpoint = torch.load(name)
+model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
+model.load_state_dict(checkpoint['model_state_dict'])
+model.eval()
+```
+## Model Architecture
+{str(self.model)} """
+        rp(model_card)
+    try:
+        # Save and upload README
+        with open("README.md", "w") as f:
+            f.write(f"# Pix2Pix UNet Model\n\n"
+                    f"- **Image Size:** {self.save_dict['model_config']['img_size']}\n"
+                    f"- **Model Type:** {'big' if big else 'small'}_UNet ({self.save_dict['model_config']['img_size']})\n"
+                    f"## Model Architecture\n{str(self.model)}")
+        
+        self.api.upload_file(
+            path_or_fileobj="README.md",
+            path_in_repo="README.md",
+            repo_id=self.repo_id,
+            token=self.token,
+            repo_type="model"
+        )
+        
+        # Clean up local files
+        os.remove(pth_name)
+        os.remove("README.md")
+        
+        print(f"Model successfully uploaded to {self.repo_id}")
+        
+    except Exception as e:
+        print(f"Error uploading model: {e}")
+
+def prepare_input(image, device='cpu'):
+    """Prepare image for inference"""
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    input_tensor = transform(image).unsqueeze(0).to(device)
+    return input_tensor
+
+def run_inference(image):
+    """Run inference on a single image"""
+    global global_model
+    if global_model is None:
+        return "Error: Model not loaded"
+    
+    global_model.eval()
+    input_tensor = prepare_input(image, device)
+    
+    with torch.no_grad():
+        output = global_model(input_tensor)
+    
+    # Convert output to image
+    output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
+    output = ((output - output.min()) / (output.max() - output.min()) * 255).astype(np.uint8)
+    rp(output[0])
+    return output
+    
+def to_hub(model, epoch, loss):
+    wrapper = UNetWrapper(model, model_repo_id, epoch, loss)
+    wrapper.push_to_hub()
+
+    
+def train_model(epochs, save_interval=1):
+    """Training function with checkpoint saving and model uploading."""
+    global global_model
+    
+    # Load combined data CSV
+    data_path = 'combined_data.csv'
+    combined_data = pd.read_csv(data_path)
+    
+    # Define the transformation
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    # Initialize dataset and dataloader
+    dataset = Pix2PixDataset(combined_data, transform, clip_tokenizer)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    model = global_model
+    criterion = nn.L1Loss()
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)  # Example scheduler
+    wrapper = UNetWrapper(model, model_repo_id, epoch=0, loss=0.0, optimizer=optimizer, scheduler=scheduler)
+
+    output_text = []
+
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+
+        for i, (original, target, original_prompt_tokens, enhanced_prompt_tokens) in enumerate(dataloader):
+            # Move data to device
+            original, target = original.to(device), target.to(device)
+            original_prompt_tokens = original_prompt_tokens.input_ids.to(device).float()
+            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device).float()
+            
+            optimizer.zero_grad()
+
+            # Forward pass
+            output = model(target)
+            img_loss = criterion(output, original)
+            total_loss = img_loss
+            total_loss.backward()
+            optimizer.step()
+
+            running_loss += total_loss.item()
+
+            if i % 10 == 0:
+                status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {total_loss.item():.8f}"
+                print(status)
+                output_text.append(status)
+        
+        # Update the epoch and loss for checkpoint
+        wrapper.epoch = epoch + 1
+        wrapper.loss = running_loss / len(dataloader)
+        
+        # Save checkpoint at specified intervals
+        if (epoch + 1) % save_interval == 0:
+            checkpoint_path = f'big_checkpoint_epoch_{epoch+1}.pth' if big else   f'small_checkpoint_epoch_{epoch+1}.pth'
+            wrapper.save_checkpoint(checkpoint_path)
+            wrapper.push_to_hub(checkpoint_path)
+
+        scheduler.step()  # Update learning rate scheduler
+
+    global_model = model  # Update global model after training
+    return model, "\n".join(output_text)
+
+    
+def train_model_old(epochs):
+    """Training function"""
+    global global_model
+    
+    # Load combined data CSV
+    data_path = 'combined_data.csv'  # Adjust this path
+    combined_data = pd.read_csv(data_path)
+    
+    # Define the transformation
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+    
+    # Initialize the dataset and dataloader
+    dataset = Pix2PixDataset(combined_data, transform, clip_tokenizer)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    model = global_model
+    criterion = nn.L1Loss()  # L1 loss for image reconstruction
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+    output_text = []
+    
+    for epoch in range(epochs):
+        model.train()
+        for i, (original, target, original_prompt_tokens, enhanced_prompt_tokens) in enumerate(dataloader):
+            # Move images and prompt embeddings to the appropriate device (CPU or GPU)
+            original, target = original.to(device), target.to(device)
+            original_prompt_tokens = original_prompt_tokens.input_ids.to(device).float()  # Convert to float
+            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device).float()  # Convert to float
+            
+            optimizer.zero_grad()
+        
+            # Forward pass through the model
+            output = model(target)
+        
+            # Compute image reconstruction loss
+            img_loss = criterion(output, original)
+            rp(f"Image {i} Loss:{img_loss}")
+            
+            # Combine losses
+            total_loss = img_loss  # Add any other losses if necessary
+            total_loss.backward()
+        
+            # Optimizer step
+            optimizer.step()
+
+            if i % 10 == 0:
+                status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {total_loss.item():.8f}"
+                rp(status)
+                output_text.append(status)
+        
+        # Push model to Hugging Face Hub at the end of each epoch
+        to_hub(model, epoch, total_loss)
+        
+    global_model = model  # Update the global model after training
+    return model, "\n".join(output_text)
+
+def gradio_train(epochs):
+    # Gradio training interface function
+    model, training_log = train_model(int(epochs))
+    #to_hub(model)
+    return f"{training_log}\n\nModel trained for {epochs} epochs and pushed to {model_repo_id}"
+
+def gradio_inference(input_image):
+    # Gradio inference interface function
+    output_image = run_inference(input_image)  # Assuming `run_inference` returns a tuple (output_image, other_data)
+    rp(output_image)
+    # If `run_inference` returns a tuple, you should only return the image part
+    return output_image  # Ensure you're only returning the processed output image
+
+
+# Create Gradio interface with tabs
+with gr.Blocks() as app:
+    gr.Markdown("# Pix2Pix Model Training and Inference")
+    
+    with gr.Tab("Train"):
+        epochs_input = gr.Number(value=EPOCHS, label="Number of epochs")
+        train_button = gr.Button("Train")
+        training_output = gr.Textbox(label="Training Log", interactive=False)
+        train_button.click(gradio_train, inputs=[epochs_input], outputs=[training_output])
+    
+    with gr.Tab("Inference"):
+        image_input = gr.Image(type='numpy')
+        prompt_input = gr.Textbox(label="Prompt")
+        inference_button = gr.Button("Generate")
+        inference_output = gr.Image(type='numpy', label="Generated Image")
+        inference_button.click(gradio_inference, inputs=[image_input], outputs=[inference_output])
+
+load_model()
+app.launch()