app.py

import spaces
import os
import torch
import random
from huggingface_hub import snapshot_download
from diffusers import StableDiffusionXLPipeline, AutoencoderKL
from diffusers import EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DPMSolverSDEScheduler
import gradio as gr
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM, pipeline

import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# Download the model files
ckpt_dir = snapshot_download(repo_id="John6666/pony-realism-v21main-sdxl")

# Load the models
vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.float16)

pipe = StableDiffusionXLPipeline.from_pretrained(
    ckpt_dir,
    vae=vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
)
pipe = pipe.to("cuda")

# Define samplers
samplers = {
    "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config),
    "DPM++ 2M": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, algorithm_type="dpmsolver++", use_karras_sigmas=True),
    "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
}

DEFAULT_POSITIVE_PREFIX = "score_9, score_8_up, score_7_up, BREAK,"
DEFAULT_POSITIVE_SUFFIX = "(masterpiece), best quality, very aesthetic, perfect face"
DEFAULT_NEGATIVE_PREFIX = "score_1, score_2, score_3, text"
DEFAULT_NEGATIVE_SUFFIX = "nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"

# Initialize Florence model
device = "cuda" if torch.cuda.is_available() else "cpu"
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)

# 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="<DETAILED_CAPTION>", 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="<DETAILED_CAPTION>",
        image_size=(image.width, image.height)
    )
    return parsed_answer["<DETAILED_CAPTION>"]

# 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']
    else:  # Long
        result = enhancer_long("Enhance the description: " + input_prompt)
        enhanced_text = result[0]['summary_text']
    
    return enhanced_text

@spaces.GPU(duration=120)
def generate_image(additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps, guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer, input_image=None, progress=gr.Progress(track_tqdm=True)):
    if use_random_seed:
        seed = random.randint(0, 2**32 - 1)
    else:
        seed = int(seed)  # Ensure seed is an integer
    
    # Set the scheduler based on the selected sampler
    pipe.scheduler = samplers[sampler]
    
    # Set clip skip
    pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)
    
    # Start with the default positive prompt prefix
    full_positive_prompt = DEFAULT_POSITIVE_PREFIX

    # Add Florence-2 caption if enabled and image is provided
    if use_florence2 and input_image is not None:
        florence2_caption = florence_caption(input_image)
        florence2_caption = florence2_caption.lower().replace('.', ',')
        additional_positive_prompt = f"{florence2_caption}, {additional_positive_prompt}" if additional_positive_prompt else florence2_caption

    # Enhance only the additional positive prompt if enhancers are enabled
    if additional_positive_prompt:
        enhanced_prompt = additional_positive_prompt
        if use_medium_enhancer:
            medium_enhanced = enhance_prompt(enhanced_prompt, "Medium")
            medium_enhanced = medium_enhanced.lower().replace('.', ',')
            enhanced_prompt = f"{enhanced_prompt}, {medium_enhanced}"
        if use_long_enhancer:
            long_enhanced = enhance_prompt(enhanced_prompt, "Long")
            long_enhanced = long_enhanced.lower().replace('.', ',')
            enhanced_prompt = f"{enhanced_prompt}, {long_enhanced}"
        full_positive_prompt += f"{enhanced_prompt}"

    # Add the default positive suffix
    full_positive_prompt += f", {DEFAULT_POSITIVE_SUFFIX}"
    
    # Combine default negative prompt with additional negative prompt
    full_negative_prompt = f"{DEFAULT_NEGATIVE_PREFIX}, {additional_negative_prompt}, {DEFAULT_NEGATIVE_SUFFIX}" if additional_negative_prompt else f"{DEFAULT_NEGATIVE_PREFIX}, {DEFAULT_NEGATIVE_SUFFIX}"
    
    try:
        image = pipe(
            prompt=full_positive_prompt,
            negative_prompt=full_negative_prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            num_images_per_prompt=num_images_per_prompt,
            generator=torch.Generator(pipe.device).manual_seed(seed)
        ).images
        return image, seed, full_positive_prompt
    except Exception as e:
        print(f"Error during image generation: {str(e)}")
        return None, seed, full_positive_prompt

# Gradio interface
with gr.Blocks(theme='bethecloud/storj_theme') as demo:
    gr.HTML("""
    <h1 align="center">Pony Realism v21 SDXL - Text-to-Image Generation</h1>
    <p align="center">
    <a href="https://huggingface.co/John6666/pony-realism-v21main-sdxl/" target="_blank">[HF Model Page]</a>
    <a href="https://civitai.com/models/372465/pony-realism" target="_blank">[civitai Model Page]</a>
    <a href="https://huggingface.co/microsoft/Florence-2-base" target="_blank">[Florence-2 Model]</a>
    <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance-Long" target="_blank">[Prompt Enhancer Long]</a>
    <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance" target="_blank">[Prompt Enhancer Medium]</a>
    </p>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            positive_prompt = gr.Textbox(label="Positive Prompt", placeholder="Add your positive prompt here")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Add your negative prompt here")
            
            with gr.Accordion("Advanced settings", open=False):
                height = gr.Slider(512, 2048, 1024, step=64, label="Height")
                width = gr.Slider(512, 2048, 1024, step=64, label="Width")
                num_inference_steps = gr.Slider(20, 50, 30, step=1, label="Number of Inference Steps")
                guidance_scale = gr.Slider(1, 20, 6, step=0.1, label="Guidance Scale")
                num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt")
                use_random_seed = gr.Checkbox(label="Use Random Seed", value=True)
                seed = gr.Number(label="Seed", value=0, precision=0)
                sampler = gr.Dropdown(label="Sampler", choices=list(samplers.keys()), value="DPM++ SDE Karras")
                clip_skip = gr.Slider(1, 4, 2, step=1, label="Clip skip")
            
            with gr.Accordion("Captioner and Enhancers", open=False):
                input_image = gr.Image(label="Input Image for Florence-2 Captioner")
                use_florence2 = gr.Checkbox(label="Use Florence-2 Captioner", value=False)
                use_medium_enhancer = gr.Checkbox(label="Use Medium Prompt Enhancer", value=False)
                use_long_enhancer = gr.Checkbox(label="Use Long Prompt Enhancer", value=False)
            
            generate_btn = gr.Button("Generate Image")

        with gr.Column(scale=1):
            output_gallery = gr.Gallery(label="Result", elem_id="gallery", show_label=False)
            seed_used = gr.Number(label="Seed Used")
            full_prompt_used = gr.Textbox(label="Full Positive Prompt Used")

    generate_btn.click(
        fn=generate_image,
        inputs=[
            positive_prompt, negative_prompt, height, width, num_inference_steps,
            guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
            clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer, input_image
        ],
        outputs=[output_gallery, seed_used, full_prompt_used]
    )

demo.launch(debug=True)