diff --git a/README.md b/README.md new file mode 100644 index 0000000..082ee3c --- /dev/null +++ b/README.md @@ -0,0 +1,27 @@ +![CLIPtion diagram](media/CLIPtion-diagram.png) + +**CLIPtion** is a fast and small captioning extension to [OpenAI CLIP](https://github.com/openai/CLIP/) ViT-L/14. You already have ViT-L loaded when using using Stable Diffusion, SDXL, SD3, FLUX, etc and with just an extra *100MB memory* you can include caption/prompt generation in your workflows! + + +# Example workflow + +... + +# Nodes + +## CLIPtion Loader + +The loader will automatically download the CLIPtion model for you the first time it is run from the [HuggingFace CLIPtion repo](https://huggingface.co/pharmapsychotic/CLIPtion). It gets stored in the [HuggingFace cache dir](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache) (controlled by `HF_HOME` environment variable). You could also manually download the safetensor file and put in comfy-cliption directory and that will be used instead. + +## CLIPtion Generate + +Create caption from an image or batch of images. +* `temperature` - controls randomness in generation - higher values produce more diverse outputs, lower values are more focused and predictable +* `best_of` - generates this many captions in parallel and picks the one with best CLIP similarity to the image +* `ramble` - forces generation of full 77 tokens + +## CLIPtion Beam Search + +* `beam_width` - how many alternative captions are considered in parallel - higher values explore more possibilities but take longer +* `ramble` - forces generation of full 77 tokens + diff --git a/__init__.py b/__init__.py index 594bc12..5f8e1d1 100644 --- a/__init__.py +++ b/__init__.py @@ -194,7 +194,6 @@ class CLIPtionModel(nn.Module): def generate_beam( self, images: torch.Tensor, - temperature: float = 0.7, beam_width: int = 4, ramble: bool = False, ) -> List[str]: @@ -214,7 +213,6 @@ class CLIPtionModel(nn.Module): image_embeds[image_idx : image_idx + 1], device, beam_width=beam_width, - temperature=temperature, ramble=ramble, ) # pick highest scoring candidate @@ -304,7 +302,6 @@ class CLIPtionModel(nn.Module): image_embed: torch.Tensor, device: torch.device, beam_width: int = 5, - temperature: float = 0.7, ramble: bool = False, ): tokenizer = self.tokenizer @@ -340,7 +337,6 @@ class CLIPtionModel(nn.Module): # get next token logits and log probabilities logits = self.output_projection(x[:, -1:]) - logits = logits / temperature if ramble: logits[:, :, tokenizer.eos_token_id] = -float("inf") log_probs = F.log_softmax(logits, dim=-1) @@ -525,10 +521,6 @@ class CLIPtionBeamSearch: ), }, "optional": { - "temperature": ( - "FLOAT", - {"default": 0.7, "tooltip": "Temperature for token sampling."}, - ), "ramble": ("BOOLEAN", {"default": False}), }, } @@ -538,11 +530,10 @@ class CLIPtionBeamSearch: model: CLIPtionModel, image: torch.Tensor, beam_width: int = 4, - temperature: float = 0.7, ramble: bool = False, ): with torch.inference_mode(): - captions = model.generate_beam(image, temperature, beam_width, ramble) + captions = model.generate_beam(image, beam_width, ramble) return (captions,) @@ -551,3 +542,9 @@ NODE_CLASS_MAPPINGS = { "CLIPtionGenerate": CLIPtionGenerate, "CLIPtionLoader": CLIPtionLoader, } + +NODE_DISPLAY_NAME_MAPPINGS = { + "CLIPtionBeamSearch": "CLIPtion Beam Search", + "CLIPtionGenerate": "CLIPtion Generate", + "CLIPtionLoader": "CLIPtion Loader", +} \ No newline at end of file diff --git a/media/CLIPtion-diagram.png b/media/CLIPtion-diagram.png new file mode 100644 index 0000000..ff55c88 Binary files /dev/null and b/media/CLIPtion-diagram.png differ