Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import gradio as gr | |
| import torch | |
| from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, pipeline | |
| from diffusers import StableDiffusion3Pipeline, StableDiffusion3Img2ImgPipeline | |
| import re | |
| import random | |
| import numpy as np | |
| import os | |
| from huggingface_hub import snapshot_download | |
| # Initialize models | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.float16 | |
| huggingface_token = os.getenv("HUGGINGFACE_TOKEN") | |
| model_path = snapshot_download( | |
| repo_id="stabilityai/stable-diffusion-3-medium", | |
| revision="refs/pr/26", | |
| repo_type="model", | |
| ignore_patterns=["*.md", "*..gitattributes"], | |
| local_dir="SD3", | |
| token=huggingface_token, # type a new token-id. | |
| ) | |
| # VLM Captioner | |
| vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to(device).eval() | |
| vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner") | |
| # 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) | |
| def load_pipeline(pipeline_type): | |
| if pipeline_type == "text2img": | |
| return StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=dtype).to(device) | |
| elif pipeline_type == "img2img": | |
| return StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=dtype).to(device) | |
| MAX_SEED = np.iinfo(np.int32).max | |
| MAX_IMAGE_SIZE = 1344 | |
| # 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, 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 | |
| # SD3 Generation function | |
| def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator().manual_seed(seed) | |
| pipe = load_pipeline("text2img") | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, | |
| width=width, | |
| height=height, | |
| generator=generator | |
| ).images[0] | |
| return image, seed | |
| # Gradio Interface | |
| def process_workflow(image, text_prompt, use_vlm, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps): | |
| if use_vlm and image is not None: | |
| prompt = create_captions_rich(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) | |
| return generated_image, prompt, used_seed | |
| def img2img_generate( | |
| prompt: str, | |
| init_image: gr.Image, | |
| use_vlm: bool, | |
| use_enhancer: bool, | |
| model_choice: str, | |
| negative_prompt: str = "", | |
| seed: int = 0, | |
| randomize_seed: bool = False, | |
| guidance_scale: float = 7, | |
| num_inference_steps: int = 30, | |
| strength: float = 0.8, | |
| ): | |
| if use_vlm and init_image is not None: | |
| prompt = create_captions_rich(init_image) | |
| if use_enhancer: | |
| prompt = enhance_prompt(prompt, model_choice) | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator().manual_seed(seed) | |
| img2img_pipe = load_pipeline("img2img") | |
| init_image = init_image.resize((768, 768)) | |
| image = img2img_pipe( | |
| prompt=prompt, | |
| image=init_image, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, | |
| generator=generator, | |
| strength=strength, | |
| ).images[0] | |
| return image, prompt, 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; | |
| } | |
| """ | |
| # Gradio Interface | |
| with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo: | |
| gr.Markdown("# VLM Captioner + Prompt Enhancer + SD3 Image Generator") | |
| with gr.Tab(label="Text to Image"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Group(elem_classes="input-group"): | |
| input_image = gr.Image(label="Input Image for VLM") | |
| use_vlm = gr.Checkbox(label="Use VLM Captioner", value=False) | |
| with gr.Group(elem_classes="input-group"): | |
| text_prompt = gr.Textbox(label="Text Prompt") | |
| use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False) | |
| model_choice = gr.Radio(["Medium", "Long"], label="Enhancer Model", value="Long") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| 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=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024) | |
| height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024) | |
| guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=5.0) | |
| num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=28) | |
| 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.Image(label="Generated Image") | |
| 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, use_vlm, use_enhancer, model_choice, | |
| negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps | |
| ], | |
| outputs=[output_image, final_prompt, used_seed] | |
| ) | |
| with gr.Tab(label="Image to Image"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Group(elem_classes="input-group"): | |
| init_image = gr.Image(label="Input Image", type="pil") | |
| use_vlm = gr.Checkbox(label="Use VLM Captioner", value=False) | |
| with gr.Group(elem_classes="input-group"): | |
| img2img_prompt = gr.Textbox(label="Text Prompt") | |
| use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False) | |
| model_choice = gr.Radio(["Medium", "Long"], label="Enhancer Model", value="Long") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| 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) | |
| guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=10.0, step=0.1, value=5) | |
| num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=28) | |
| strength = gr.Slider(label="Img2Img Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.5) | |
| img2img_generate_btn = gr.Button("Generate Image", elem_classes="submit-btn") | |
| with gr.Column(scale=1): | |
| with gr.Group(elem_classes="output-group"): | |
| img2img_output = gr.Image(label="Generated Image") | |
| img2img_final_prompt = gr.Textbox(label="Final Prompt Used") | |
| img2img_used_seed = gr.Number(label="Seed Used") | |
| img2img_generate_btn.click( | |
| fn=img2img_generate, | |
| inputs=[ | |
| img2img_prompt, init_image, use_vlm, use_enhancer, model_choice, | |
| negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, strength | |
| ], | |
| outputs=[img2img_output, img2img_final_prompt, img2img_used_seed] | |
| ) | |
| demo.launch(debug=True) |