import spaces import gradio as gr import torch from PIL import Image from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, pipeline from transformers import AutoProcessor, AutoModelForCausalLM from transformers.dynamic_module_utils import get_imports from unittest.mock import patch import re import random import os import time os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60") os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60") from huggingface_hub import snapshot_download from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline from kolors.models.modeling_chatglm import ChatGLMModel from kolors.models.tokenization_chatglm import ChatGLMTokenizer from diffusers import UNet2DConditionModel, AutoencoderKL from diffusers import EulerDiscreteScheduler # Initialize models device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float16 def snapshot_download_with_retries(repo_id, retries=3): last_error = None for attempt in range(1, retries + 1): try: return snapshot_download(repo_id=repo_id, etag_timeout=60, max_workers=2) except Exception as exc: last_error = exc if attempt == retries: raise print(f"Snapshot download failed for {repo_id} on attempt {attempt}: {exc}") time.sleep(5 * attempt) raise last_error # Download Kolors model ckpt_dir = snapshot_download_with_retries(repo_id="Kwai-Kolors/Kolors") def fixed_get_imports(filename): if not str(filename).endswith("/modeling_florence2.py"): return get_imports(filename) imports = get_imports(filename) for optional_import in ("flash_attn", "einops"): if optional_import == "flash_attn" and optional_import in imports: imports.remove(optional_import) return imports # Load Kolors models text_encoder = ChatGLMModel.from_pretrained(os.path.join(ckpt_dir, 'text_encoder'), torch_dtype=dtype).to(device) tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(ckpt_dir, 'text_encoder')) vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), revision=None).to(dtype).to(device) scheduler = EulerDiscreteScheduler.from_pretrained(os.path.join(ckpt_dir, "scheduler")) unet = UNet2DConditionModel.from_pretrained(os.path.join(ckpt_dir, "unet"), revision=None).to(dtype).to(device) kolors_pipe = StableDiffusionXLPipeline( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler, force_zeros_for_empty_prompt=False ).to(device) # VLM Captioner vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner-v2").to(device).eval() vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner-v2") # Initialize Florence model with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval() florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True) # Prompt Enhancer enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device) enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device) MAX_SEED = 2**32 - 1 # Florence caption function def florence_caption(image): # Convert image to PIL if it's not already if not isinstance(image, Image.Image): image = Image.fromarray(image) inputs = florence_processor(text="", images=image, return_tensors="pt").to(device) generated_ids = florence_model.generate( input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, early_stopping=False, do_sample=False, num_beams=3, ) generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0] parsed_answer = florence_processor.post_process_generation( generated_text, task="", image_size=(image.width, image.height) ) return parsed_answer[""] # VLM Captioner function def create_captions_rich(image): prompt = "caption en" model_inputs = vlm_processor(text=prompt, images=image, return_tensors="pt").to(device) input_len = model_inputs["input_ids"].shape[-1] with torch.inference_mode(): generation = vlm_model.generate(**model_inputs, repetition_penalty=1.10, max_new_tokens=256, do_sample=False) generation = generation[0][input_len:] decoded = vlm_processor.decode(generation, skip_special_tokens=True) return modify_caption(decoded) # Helper function for caption modification def modify_caption(caption: str) -> str: prefix_substrings = [ ('captured from ', ''), ('captured at ', '') ] pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings]) replacers = {opening: replacer for opening, replacer in prefix_substrings} def replace_fn(match): return replacers[match.group(0)] return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE) # Prompt Enhancer function def enhance_prompt(input_prompt, model_choice): if model_choice == "Medium": result = enhancer_medium("Enhance the description: " + input_prompt) enhanced_text = result[0]['summary_text'] pattern = r'^.*?of\s+(.*?(?:\.|$))' match = re.match(pattern, enhanced_text, re.IGNORECASE | re.DOTALL) if match: remaining_text = enhanced_text[match.end():].strip() modified_sentence = match.group(1).capitalize() enhanced_text = modified_sentence + ' ' + remaining_text else: # Long result = enhancer_long("Enhance the description: " + input_prompt) enhanced_text = result[0]['summary_text'] return enhanced_text def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt): if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) image = kolors_pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, num_images_per_prompt=num_images_per_prompt, generator=generator ).images return image, seed @spaces.GPU(duration=200) def process_workflow(image, text_prompt, vlm_model_choice, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt): if image is not None: # Convert image to PIL if it's not already if not isinstance(image, Image.Image): image = Image.fromarray(image) if vlm_model_choice == "Long Captioner": prompt = create_captions_rich(image) else: # Florence prompt = florence_caption(image) else: prompt = text_prompt if use_enhancer: prompt = enhance_prompt(prompt, model_choice) generated_image, used_seed = generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt) return generated_image, prompt, used_seed custom_css = """ .input-group, .output-group { border: 1px solid #e0e0e0; border-radius: 10px; padding: 20px; margin-bottom: 20px; background-color: #f9f9f9; } .submit-btn { background-color: #2980b9 !important; color: white !important; } .submit-btn:hover { background-color: #3498db !important; } """ title = """

Kolors with VLM Captioner and Prompt Enhancer

[Kolors Model] [Florence-2 Model] [Long Captioner Model] [Prompt Enhancer Long] [Prompt Enhancer Medium]

Create long prompts from images or enhance your short prompts with prompt enhancer

""" # gokaygokay-ui-citation CITATION_MARKDOWN = 'This hosted interface is maintained by [Gökay Aydoğan](https://orcid.org/0000-0002-2343-9433). If you reference this repository in academic work, please cite it as follows and also cite the upstream models, datasets, or projects it builds upon.\n\n```bibtex\n@software{aydogan2024kolorsplusplus,\n author = {Aydoğan, Gökay},\n title = {{KolorsPlusPlus}},\n year = {2024},\n publisher = {Hugging Face},\n url = {https://huggingface.co/spaces/gokaygokay/KolorsPlusPlus},\n note = {Hosted interface; cite the upstream models and projects used by this Space as required.}\n}\n```' with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo: with gr.Accordion("Citation & attribution", open=False): gr.Markdown(CITATION_MARKDOWN) gr.HTML(title) with gr.Row(): with gr.Column(scale=1): with gr.Group(elem_classes="input-group"): input_image = gr.Image(label="Input Image (VLM Captioner)") vlm_model_choice = gr.Radio(["Florence-2", "Long Captioner"], label="VLM Model", value="Florence-2") with gr.Accordion("Advanced Settings", open=False): text_prompt = gr.Textbox(label="Text Prompt (optional, used if no image is uploaded)") use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False) model_choice = gr.Radio(["Medium", "Long"], label="Enhancer Model", value="Long") negative_prompt = gr.Textbox(label="Negative Prompt") seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) width = gr.Slider(label="Width", minimum=512, maximum=2048, step=64, value=1024) height = gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1024) guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=5.0) num_inference_steps = gr.Slider(label="Inference Steps", minimum=20, maximum=50, step=1, value=20) num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt") generate_btn = gr.Button("Generate Image", elem_classes="submit-btn") with gr.Column(scale=1): with gr.Group(elem_classes="output-group"): output_image = gr.Gallery(label="Result", elem_id="gallery", show_label=False) final_prompt = gr.Textbox(label="Final Prompt Used") used_seed = gr.Number(label="Seed Used") generate_btn.click( fn=process_workflow, inputs=[ input_image, text_prompt, vlm_model_choice, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt ], outputs=[output_image, final_prompt, used_seed] ) demo.launch(debug=True)