Fix for multiple graphics cards

README.md

@@ -2,7 +2,7 @@
 
 ## Overview
 
-This application generates descriptive captions for images using advanced ML models. It processes single images or entire directories, leveraging CLIP and LLM models for accurate and contextual captions. It has NSFW captioning support with natural language. This is just an extension of the original author's efforts to improve performance. Their repo is located here: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha.
+This application generates descriptive captions for images using advanced ML models. It processes single images or entire directories, leveraging CLIP and LLM models for accurate and contextual captions. It has NSFW captioning support with natural language. This is just an extension of the original author's efforts to improve performance. Their report is located here: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha.
 
 ## Features
 
@@ -44,8 +44,6 @@ git clone https://huggingface.co/Wi-zz/joy-caption-pre-alpha
 cd joy-caption-pre-alpha
 python -m venv venv
 .\venv\Scripts\activate
-# Change as per https://pytorch.org/get-started/locally/
-pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
 pip install -r requirements.txt
 ```
 
@@ -56,7 +54,6 @@ git clone https://huggingface.co/Wi-zz/joy-caption-pre-alpha
 cd joy-caption-pre-alpha
 python3 -m venv venv
 source venv/bin/activate
-pip3 install torch torchvision torchaudio
 pip3 install -r requirements.txt
 ```
 

app-multi-alpha.py

@@ -1,210 +0,0 @@
-import torch
-import torch.amp.autocast_mode
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from torch.nn.parallel import DistributedDataParallel as DDP
-import os
-import sys
-import logging
-import warnings
-import argparse
-from PIL import Image
-from pathlib import Path
-from tqdm import tqdm
-from torch import nn
-from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
-from typing import List, Union
-
-# Constants
-CLIP_PATH = "google/siglip-so400m-patch14-384"
-VLM_PROMPT = "A descriptive caption for this image:\n"
-MODEL_PATH = "unsloth/Meta-Llama-3.1-8B-bnb-4bit"
-CHECKPOINT_PATH = Path("wpkklhc6")
-IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp', '.webp')
-
-warnings.filterwarnings("ignore", category=UserWarning)
-logging.getLogger("transformers").setLevel(logging.ERROR)
-
-def setup(rank, world_size):
-    os.environ['MASTER_ADDR'] = 'localhost'
-    os.environ['MASTER_PORT'] = '12355'
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
-
-def cleanup():
-    dist.destroy_process_group()
-
-class ImageAdapter(nn.Module):
-    def __init__(self, input_features: int, output_features: int):
-        super().__init__()
-        self.linear1 = nn.Linear(input_features, output_features)
-        self.activation = nn.GELU()
-        self.linear2 = nn.Linear(output_features, output_features)
-    
-    def forward(self, vision_outputs: torch.Tensor):
-        return self.linear2(self.activation(self.linear1(vision_outputs)))
-
-def load_models(rank):
-    print(f"Loading CLIP 📎 on GPU {rank}")
-    clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
-    clip_model = AutoModel.from_pretrained(CLIP_PATH).vision_model.eval().requires_grad_(False).to(rank)
-
-    print(f"Loading tokenizer 🪙 on GPU {rank}")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
-    assert isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)), f"Tokenizer is of type {type(tokenizer)}"
-
-    print(f"Loading LLM 🤖 on GPU {rank}")
-    text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map={"": rank}, torch_dtype=torch.bfloat16).eval()
-
-    print(f"Loading image adapter 🖼️ on GPU {rank}")
-    image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size)
-    image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location=f"cuda:{rank}", weights_only=True))
-    image_adapter.eval().to(rank)
-
-    return clip_processor, clip_model, tokenizer, text_model, image_adapter
-
-@torch.no_grad()
-def stream_chat(input_images: List[Image.Image], batch_size: int, pbar: tqdm, models: tuple, rank: int) -> List[str]:
-    clip_processor, clip_model, tokenizer, text_model, image_adapter = models
-    torch.cuda.empty_cache()
-    all_captions = []
-
-    for i in range(0, len(input_images), batch_size):
-        batch = input_images[i:i+batch_size]
-        
-        try:
-            images = clip_processor(images=batch, return_tensors='pt', padding=True).pixel_values.to(rank)
-        except ValueError as e:
-            print(f"Error processing image batch: {e}")
-            print("Skipping this batch and continuing...")
-            continue
-
-        with torch.amp.autocast_mode.autocast(device_type='cuda', enabled=True):
-            vision_outputs = clip_model(pixel_values=images, output_hidden_states=True)
-            image_features = vision_outputs.hidden_states[-2]
-            embedded_images = image_adapter(image_features).to(dtype=torch.bfloat16)
-
-        prompt = tokenizer.encode(VLM_PROMPT, return_tensors='pt')
-        prompt_embeds = text_model.model.embed_tokens(prompt.to(rank)).to(dtype=torch.bfloat16)
-        embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=rank, dtype=torch.int64)).to(dtype=torch.bfloat16)
-
-        inputs_embeds = torch.cat([
-            embedded_bos.expand(embedded_images.shape[0], -1, -1),
-            embedded_images,
-            prompt_embeds.expand(embedded_images.shape[0], -1, -1),
-        ], dim=1).to(dtype=torch.bfloat16)
-
-        input_ids = torch.cat([
-            torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).expand(embedded_images.shape[0], -1),
-            torch.zeros((embedded_images.shape[0], embedded_images.shape[1]), dtype=torch.long),
-            prompt.expand(embedded_images.shape[0], -1),
-        ], dim=1).to(rank)
-
-        attention_mask = torch.ones_like(input_ids)
-
-        generate_ids = text_model.generate(
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            max_new_tokens=300,
-            do_sample=True,
-            top_k=10,
-            temperature=0.5,
-        )
-
-        generate_ids = generate_ids[:, input_ids.shape[1]:]
-
-        for ids in generate_ids:
-            caption = tokenizer.decode(ids[:-1] if ids[-1] == tokenizer.eos_token_id else ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-            caption = caption.replace('<|end_of_text|>', '').replace('<|finetune_right_pad_id|>', '').strip()
-            all_captions.append(caption)
-
-        if pbar and rank == 0:
-            pbar.update(len(batch))
-
-    return all_captions
-
-def process_directory(rank, world_size, input_dir: Path, output_dir: Path, batch_size: int, models: tuple):
-    output_dir.mkdir(parents=True, exist_ok=True)
-    image_files = [f for f in input_dir.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS]
-    images_to_process = [f for f in image_files if not (output_dir / f"{f.stem}.txt").exists()]
-
-    if not images_to_process:
-        if rank == 0:
-            print("No new images to process.")
-        return
-
-    # Distribute images across GPUs
-    images_per_gpu = len(images_to_process) // world_size
-    start_idx = rank * images_per_gpu
-    end_idx = start_idx + images_per_gpu if rank < world_size - 1 else len(images_to_process)
-    gpu_images = images_to_process[start_idx:end_idx]
-
-    if rank == 0:
-        pbar = tqdm(total=len(images_to_process), desc="Processing images", unit="image")
-    else:
-        pbar = None
-
-    for i in range(0, len(gpu_images), batch_size):
-        batch_files = gpu_images[i:i+batch_size]
-        batch_images = [Image.open(f).convert('RGB') for f in batch_files]
-
-        captions = stream_chat(batch_images, batch_size, pbar, models, rank)
-        
-        for file, caption in zip(batch_files, captions):
-            with open(output_dir / f"{file.stem}.txt", 'w', encoding='utf-8') as f:
-                f.write(caption)
-
-        for img in batch_images:
-            img.close()
-
-    if rank == 0:
-        pbar.close()
-
-def parse_arguments():
-    parser = argparse.ArgumentParser(description="Process images and generate captions.")
-    parser.add_argument("input", nargs='+', help="Input image file or directory (or multiple directories)")
-    parser.add_argument("--output", help="Output directory (optional)")
-    parser.add_argument("--bs", type=int, default=4, help="Batch size (default: 4)")
-    return parser.parse_args()
-
-def run(rank, world_size, args):
-    setup(rank, world_size)
-    
-    input_paths = [Path(input_path) for input_path in args.input]
-    batch_size = args.bs
-    models = load_models(rank)
-
-    for input_path in input_paths:
-        if input_path.is_file() and input_path.suffix.lower() in IMAGE_EXTENSIONS:
-            if rank == 0:
-                output_path = input_path.with_suffix('.txt')
-                print(f"Processing single image 🎞️: {input_path.name}")
-                with tqdm(total=1, desc="Processing image", unit="image") as pbar:
-                    captions = stream_chat([Image.open(input_path).convert('RGB')], 1, pbar, models, rank)
-                    with open(output_path, 'w', encoding='utf-8') as f:
-                        f.write(captions[0])
-                print(f"Output saved to {output_path}")
-        elif input_path.is_dir():
-            output_path = Path(args.output) if args.output else input_path
-            if rank == 0:
-                print(f"Processing directory 📁: {input_path}")
-                print(f"Output directory 📦: {output_path}")
-                print(f"Batch size 🗄️: {batch_size}")
-            process_directory(rank, world_size, input_path, output_path, batch_size, models)
-        else:
-            if rank == 0:
-                print(f"Invalid input: {input_path}")
-                print("Skipping...")
-
-    cleanup()
-
-def main():
-    args = parse_arguments()
-    world_size = torch.cuda.device_count()
-    if world_size > 1:
-        mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True)
-    else:
-        run(0, 1, args)
-
-if __name__ == "__main__":
-    main()