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 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 # Device configuration device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') big = True if device == torch.device('cpu') else False # Parameters IMG_SIZE = 1024 if big else 256 BATCH_SIZE = 1 if big else 4 EPOCHS = 12 LR = 0.0002 dataset_id = "K00B404/pix2pix_flux_set" model_repo_id = "K00B404/pix2pix_flux" # Global model variable global_model = None def load_model(): """Load the model 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 print("Model loaded successfully!") return model except Exception as e: print(f"Error loading model: {e}") model = big_UNet().to(device) if big else small_UNet().to(device) global_model = model return model # Dataset class remains the same class Pix2PixDataset(torch.utils.data.Dataset): def __init__(self, ds, transform): self.originals = [x for x in ds["train"] if x['label'] == 0] self.targets = [x for x in ds["train"] if x['label'] == 1] assert len(self.originals) == len(self.targets) print(f"Number of original images: {len(self.originals)}") print(f"Number of target images: {len(self.targets)}") self.transform = transform def __len__(self): return len(self.originals) def __getitem__(self, idx): original_img = self.originals[idx]['image'] target_img = self.targets[idx]['image'] original = original_img.convert('RGB') target = target_img.convert('RGB') return self.transform(original), self.transform(target) class UNetWrapper: def __init__(self, unet_model, repo_id): self.model = unet_model self.repo_id = repo_id self.token = os.getenv('NEW_TOKEN') # Make sure this environment variable is set self.api = HfApi(token=os.getenv('NEW_TOKEN')) def push_to_hub(self): try: # Save model state and configuration save_dict = { 'model_state_dict': self.model.state_dict(), 'model_config': { 'big': isinstance(self.model, big_UNet), 'img_size': 1024 if isinstance(self.model, big_UNet) else 256 }, 'model_architecture': str(self.model) } # Save model locally pth_name = 'big_model_weights.pth' if big else 'small_model_weights.pth' torch.save(save_dict, pth_name) # Create repo if it doesn't exist try: create_repo( repo_id=self.repo_id, token=self.token, exist_ok=True ) except Exception as e: print(f"Repository creation note: {e}") # Upload the model file 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" ) # 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:** 1024 - **Model Type:** Big (1024) ## 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)} """ # Save and upload README with open("README.md", "w") as f: f.write(model_card) 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) return output def to_hub(model): wrapper = UNetWrapper(model, model_repo_id) wrapper.push_to_hub() def train_model(epochs): """Training function""" global global_model ds = load_dataset(dataset_id) transform = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), ]) dataset = Pix2PixDataset(ds, transform) dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True) model = global_model criterion = nn.L1Loss() optimizer = optim.Adam(model.parameters(), lr=LR) output_text = [] for epoch in range(epochs): model.train() for i, (original, target) in enumerate(dataloader): original, target = original.to(device), target.to(device) optimizer.zero_grad() output = model(target) loss = criterion(output, original) loss.backward() optimizer.step() if i % 10 == 0: status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}" print(status) output_text.append(status) to_hub(model) global_model = model 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""" return input_image, run_inference(input_image) # Create Gradio interface with tabs with gr.Blocks() as app: gr.Markdown("# Pix2Pix Model Training and Inference") with gr.Tabs(): with gr.TabItem("Training"): epochs_input = gr.Number(label="Number of Epochs") train_button = gr.Button("Train Model") output_text = gr.Textbox(label="Training Progress", lines=10) train_button.click(gradio_train, inputs=epochs_input, outputs=output_text) with gr.TabItem("Inference"): with gr.Row(): input_image = gr.Image(label="Input Image") output_image = gr.Image(label="Model Output") infer_button = gr.Button("Run Inference") infer_button.click(gradio_inference, inputs=input_image, outputs=[input_image, output_image]) if __name__ == '__main__': # Load model at startup load_model() # Launch the Gradio app app.launch()