app.py

import os
import gc
import gradio as gr
import numpy as np
import torch
import json
import spaces
import config
import utils
import logging
from PIL import Image, PngImagePlugin
from datetime import datetime
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
import random

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
HISTORY_SECRET = os.getenv("HISTORY_SECRET", "default_secret")

MODEL = os.getenv(
    "MODEL",
    "https://huggingface.co/AstraliteHeart/pony-diffusion-v6/blob/main/v6.safetensors",
)

DESCRIPTION = '''
<div>
<h1 style="text-align: center;">High Definition Pony Diffusion</h1>
<p>Gradio demo for PonyDiffusion v6 with image gallery, json prompt support, advanced options and more.</p>
<p>❤️ Thanks for ✨10k visits! FLUX INTEGRATION COMING SOON!</p>
<p>🔎 For more details about me, take a look at <a href="https://sergidev.me">My website</a>.</p>
<p>🌚 For dark mode compatibility, click <a href="https://sergidev.me/hdiffusion">here</a>.</p>
</div>
'''

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def load_pipeline(model_name):
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16,
    )
    pipeline = (
        StableDiffusionXLPipeline.from_single_file
        if MODEL.endswith(".safetensors")
        else StableDiffusionXLPipeline.from_pretrained
    )

    pipe = pipeline(
        model_name,
        vae=vae,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
        variant="fp16",
    )

    pipe.to(device)
    return pipe

def parse_json_parameters(json_str):
    try:
        params = json.loads(json_str)
        required_keys = ['prompt', 'negative_prompt', 'resolution', 'guidance_scale', 'num_inference_steps', 'seed', 'sampler']
        for key in required_keys:
            if key not in params:
                raise ValueError(f"Missing required key: {key}")
        
        width, height = map(int, params['resolution'].split(' x '))
        
        return {
            'prompt': params['prompt'],
            'negative_prompt': params['negative_prompt'],
            'seed': params['seed'],
            'width': width,
            'height': height,
            'guidance_scale': params['guidance_scale'],
            'num_inference_steps': params['num_inference_steps'],
            'sampler': params['sampler'],
            'use_upscaler': params.get('use_upscaler', False)
        }
    except json.JSONDecodeError:
        raise ValueError("Invalid JSON format")
    except Exception as e:
        raise ValueError(f"Error parsing JSON: {str(e)}")

@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    custom_width: int = 1024,
    custom_height: int = 1024,
    guidance_scale: float = 7.0,
    num_inference_steps: int = 30,
    sampler: str = "DPM++ 2M SDE Karras",
    aspect_ratio_selector: str = "1024 x 1024",
    use_upscaler: bool = False,
    upscaler_strength: float = 0.55,
    upscale_by: float = 1.5,
    json_params: str = "",
    batch_size: int = 1,
    progress=gr.Progress(track_tqdm=True),
) -> Image:
    if json_params:
        try:
            params = parse_json_parameters(json_params)
            prompt = params['prompt']
            negative_prompt = params['negative_prompt']
            seed = params['seed']
            custom_width = params['width']
            custom_height = params['height']
            guidance_scale = params['guidance_scale']
            num_inference_steps = params['num_inference_steps']
            sampler = params['sampler']
            use_upscaler = params['use_upscaler']
        except ValueError as e:
            raise gr.Error(str(e))

    generator = utils.seed_everything(seed)

    width, height = utils.aspect_ratio_handler(
        aspect_ratio_selector,
        custom_width,
        custom_height,
    )

    width, height = utils.preprocess_image_dimensions(width, height)

    backup_scheduler = pipe.scheduler
    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)

    if use_upscaler:
        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
    metadata = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": f"{width} x {height}",
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed,
        "sampler": sampler,
        "batch_size": batch_size,
    }

    if use_upscaler:
        new_width = int(width * upscale_by)
        new_height = int(height * upscale_by)
        metadata["use_upscaler"] = {
            "upscale_method": "nearest-exact",
            "upscaler_strength": upscaler_strength,
            "upscale_by": upscale_by,
            "new_resolution": f"{new_width} x {new_height}",
        }
    else:
        metadata["use_upscaler"] = None
    logger.info(json.dumps(metadata, indent=4))

    try:
        all_images = []
        for _ in range(batch_size):
            batch_generator = utils.seed_everything(random.randint(0, utils.MAX_SEED))
            if use_upscaler:
                latents = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=width,
                    height=height,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=batch_generator,
                    output_type="latent",
                ).images
                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
                images = upscaler_pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    image=upscaled_latents,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    strength=upscaler_strength,
                    generator=batch_generator,
                    output_type="pil",
                ).images
            else:
                images = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=width,
                    height=height,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=batch_generator,
                    output_type="pil",
                ).images
            all_images.extend(images)

        if all_images and IS_COLAB:
            for image in all_images:
                filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                logger.info(f"Image saved as {filepath} with metadata")

        return all_images, metadata
    except Exception as e:
        logger.exception(f"An error occurred: {e}")
        raise
    finally:
        if use_upscaler:
            del upscaler_pipe
        pipe.scheduler = backup_scheduler
        utils.free_memory()

generation_history = []

def update_history_list():
    return [item["image"] for item in generation_history]

def handle_image_click(evt: gr.SelectData):
    selected = generation_history[evt.index]
    return selected["image"], json.dumps(selected["metadata"], indent=2)

def generate_and_update_history(*args, **kwargs):
    global generation_history
    images, metadata = generate(*args, **kwargs)
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    for image in images:
        generation_history.insert(0, {
            "prompt": metadata["prompt"],
            "timestamp": timestamp,
            "image": image,
            "metadata": metadata
        })
    if len(generation_history) > 20:
        generation_history = generation_history[:20]
    return images[0], json.dumps(metadata, indent=2), update_history_list()

with open('characterfull.txt', 'r') as f:
    characters = [line.strip() for line in f.readlines()]

def get_random_character():
    return random.choice(characters)

def add_quality_tags(prompt, negative_prompt):
    positive_tags = "score_9, score_8_up, score_7_up, score_6_up, dramatic lighting"
    negative_tags = "worst quality, low quality, text, censored, deformed, bad hand, blurry, (watermark), mutated hands, monochrome"
    
    new_prompt = f"{positive_tags}, {prompt}" if prompt else positive_tags
    new_negative_prompt = f"{negative_tags}, {negative_prompt}" if negative_prompt else negative_tags
    
    return new_prompt, new_negative_prompt

if torch.cuda.is_available():
    pipe = load_pipeline(MODEL)
    logger.info("Loaded on Device!")
else:
    pipe = None

def check_history_password(password):
    if password == HISTORY_SECRET:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    )
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=5,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button(
                "Generate", 
                variant="primary", 
                scale=0
            )
        with gr.Row():
            random_character_button = gr.Button("Random Character")
            add_quality_tags_button = gr.Button("Add quality tags")
        result = gr.Image(
            label="Result", 
            show_label=False
        )
    with gr.Accordion(label="Advanced Settings", open=False):
        negative_prompt = gr.Text(
            label="Negative Prompt",
            max_lines=5,
            placeholder="Enter a negative prompt",
            value=""
        )
        aspect_ratio_selector = gr.Radio(
            label="Aspect Ratio",
            choices=config.aspect_ratios,
            value="1024 x 1024",
            container=True,
        )
        with gr.Group(visible=False) as custom_resolution:
            with gr.Row():
                custom_width = gr.Slider(
                    label="Width",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )
                custom_height = gr.Slider(
                    label="Height",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )
        use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
        with gr.Row() as upscaler_row:
            upscaler_strength = gr.Slider(
                label="Strength",
                minimum=0,
                maximum=1,
                step=0.05,
                value=0.55,
                visible=False,
            )
            upscale_by = gr.Slider(
                label="Upscale by",
                minimum=1,
                maximum=1.5,
                step=0.1,
                value=1.5,
                visible=False,
            )
        sampler = gr.Dropdown(
            label="Sampler",
            choices=config.sampler_list,
            interactive=True,
            value="DPM++ 2M SDE Karras",
        )
        with gr.Row():
            seed = gr.Slider(
                label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Group():
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1,
                    maximum=12,
                    step=0.1,
                    value=7.0,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        batch_size = gr.Slider(
            label="Batch Size",
            minimum=1,
            maximum=4,
            step=1,
            value=1,
        )
                
    with gr.Accordion(label="Generation Parameters", open=False):
        gr_metadata = gr.JSON(label="Metadata", show_label=False)
        json_input = gr.TextArea(label="Edit/Paste JSON Parameters", placeholder="Paste or edit JSON parameters here")
        generate_from_json = gr.Button("Generate from JSON")

    with gr.Accordion("Generation History", open=False) as history_accordion:
        history_password = gr.Textbox(
            label="Enable generation history; do not generate illegal or harmful content.",
            type="password",
            placeholder="GLOBAL GENERATION HISTORY IS DISABLED"
        )
        history_submit = gr.Button("Submit")
        
        with gr.Group(visible=False) as history_content:
            history_gallery = gr.Gallery(
                label="History",
                show_label=False,
                elem_id="history_gallery",
                columns=5,
                rows=2,
                height="auto"
            )
            with gr.Row():
                selected_image = gr.Image(label="Selected Image", interactive=False)
                selected_metadata = gr.JSON(label="Selected Metadata", show_label=False)

    gr.Examples(
        examples=config.examples,
        inputs=prompt,
        outputs=[result, gr_metadata],
        fn=lambda *args, **kwargs: generate_and_update_history(*args, use_upscaler=True, **kwargs),
        cache_examples=CACHE_EXAMPLES,
    )

    use_upscaler.change(
        fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
        inputs=use_upscaler,
        outputs=[upscaler_strength, upscale_by],
        queue=False,
        api_name=False,
    )
    
    aspect_ratio_selector.change(
        fn=lambda x: gr.update(visible=x == "Custom"),
        inputs=aspect_ratio_selector,
        outputs=custom_resolution,
        queue=False,
        api_name=False,
    )

    inputs = [
        prompt,
        negative_prompt,
        seed,
        custom_width,
        custom_height,
        guidance_scale,
        num_inference_steps,
        sampler,
        aspect_ratio_selector,
        use_upscaler,
        upscaler_strength,
        upscale_by,
        json_input,
        batch_size,
    ]

    prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    negative_prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    run_button.click(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    generate_from_json.click(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )

    random_character_button.click(
        fn=get_random_character,
        inputs=[],
        outputs=[prompt]
    )

    add_quality_tags_button.click(
        fn=add_quality_tags,
        inputs=[prompt, negative_prompt],
        outputs=[prompt, negative_prompt]
    )

    history_gallery.select(
        fn=handle_image_click,
        inputs=[],
        outputs=[selected_image, selected_metadata]
    )

    history_submit.click(
        fn=check_history_password,
        inputs=[history_password],
        outputs=[history_content],
    )
    
demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)