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,
    HeunDiscreteScheduler,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UniPCMultistepScheduler,
)
from diffusers.models.attention_processor import AttnProcessor2_0
import gradio as gr
from PIL import Image
import numpy as np
from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
import requests

def download_file(url, folder_path, filename):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    file_path = os.path.join(folder_path, filename)
    if os.path.isfile(file_path):
        print(f"File already exists: {file_path}")
    else:
        response = requests.get(url, stream=True)
        if response.status_code == 200:
            with open(file_path, 'wb') as file:
                for chunk in response.iter_content(chunk_size=1024):
                    file.write(chunk)
            print(f"File successfully downloaded and saved: {file_path}")
        else:
            print(f"Error downloading the file. Status code: {response.status_code}")

ckpt_dir_pony = snapshot_download(repo_id="John6666/pony-realism-v22main-sdxl")
ckpt_dir_cyber = snapshot_download(repo_id="John6666/cyberrealistic-pony-v65-sdxl")
ckpt_dir_stallion = snapshot_download(repo_id="TheImposterImposters/tamePony_v25")

vae_pony = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_pony, "vae"), torch_dtype=torch.float16)
vae_cyber = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_cyber, "vae"), torch_dtype=torch.float16)
vae_stallion = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_stallion, "vae"), torch_dtype=torch.float16)

pipe_pony = StableDiffusionXLPipeline.from_pretrained(
    ckpt_dir_pony, vae=vae_pony, torch_dtype=torch.float16, use_safetensors=True,
)
pipe_cyber = StableDiffusionXLPipeline.from_pretrained(
    ckpt_dir_cyber, vae=vae_cyber, torch_dtype=torch.float16, use_safetensors=True,
)
pipe_stallion = StableDiffusionXLPipeline.from_pretrained(
    ckpt_dir_stallion, vae=vae_stallion, torch_dtype=torch.float16, use_safetensors=True,
)

pipe_pony = pipe_pony.to("cuda")
pipe_cyber = pipe_cyber.to("cuda")
pipe_stallion = pipe_stallion.to("cuda")

pipe_pony.unet.set_attn_processor(AttnProcessor2_0())
pipe_cyber.unet.set_attn_processor(AttnProcessor2_0())
pipe_stallion.unet.set_attn_processor(AttnProcessor2_0())

samplers = {
    "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe_pony.scheduler.config),
    "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe_pony.scheduler.config, use_karras_sigmas=True),
    "Heun": HeunDiscreteScheduler.from_config(pipe_pony.scheduler.config),
    "DPM++ 2M SDE Karras": DPMSolverMultistepScheduler.from_config(pipe_pony.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"),
    "DPM++ 2M": DPMSolverMultistepScheduler.from_config(pipe_pony.scheduler.config),
    "DDIM": DDIMScheduler.from_config(pipe_pony.scheduler.config),
    "LMS": LMSDiscreteScheduler.from_config(pipe_pony.scheduler.config),
    "PNDM": PNDMScheduler.from_config(pipe_pony.scheduler.config),
    "UniPC": UniPCMultistepScheduler.from_config(pipe_pony.scheduler.config),
}

DEFAULT_POSITIVE_PREFIX = "Score_9 score_8_up score_7_up BREAK"
DEFAULT_POSITIVE_SUFFIX = "(masterpiece) very_aesthetic detailed_face cinematic footage"
DEFAULT_NEGATIVE_PREFIX = "Score_1 score_2 score _3 text low_res"
DEFAULT_NEGATIVE_SUFFIX = "Nsfw oversaturated crappy_art low_quality blurry bad_anatomy extra_digits fewer_digits simple_background very_displeasing watermark signature"

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

enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)

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

@spaces.GPU(duration=120)
def generate_image(model_choice, 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_medium_enhancer, use_long_enhancer,
                   use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
                   progress=gr.Progress(track_tqdm=True)):

    if model_choice == "Pony Realism v22":
        pipe = pipe_pony
    elif model_choice == "Cyber Realistic Pony v65":
        pipe = pipe_cyber
    else:
        pipe = pipe_stallion

    if use_random_seed:
        seed = random.randint(0, 2**32 - 1)
    else:
        seed = int(seed)

    pipe.scheduler = samplers[sampler]
    pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)

    full_positive_prompt = DEFAULT_POSITIVE_PREFIX + ", " if use_positive_prefix else ""

    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 += enhanced_prompt

    if use_positive_suffix:
        full_positive_prompt += f", {DEFAULT_POSITIVE_SUFFIX}"

    full_negative_prompt = ""
    if use_negative_prefix:
        full_negative_prompt += f"{DEFAULT_NEGATIVE_PREFIX}, "
    full_negative_prompt += additional_negative_prompt if additional_negative_prompt else ""
    if use_negative_suffix:
        full_negative_prompt += f", {DEFAULT_NEGATIVE_SUFFIX}"

    try:
        images = 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 images, seed, full_positive_prompt, full_negative_prompt
    except Exception as e:
        print(f"Error during image generation: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, seed, full_positive_prompt, full_negative_prompt

with gr.Blocks(theme='bethecloud/storj_theme') as demo:
    gr.HTML("""
<h1 align="center">Pony Realism / Cyber Realism / Stallion Dreams</h1>
<p align="center">
<a href="https://huggingface.co/John6666/pony-realism-v22main-sdxl/" target="_blank">[Pony Realism]</a>
<a href="https://huggingface.co/John6666/cyberrealistic-pony-v65-sdxl" target="_blank">[Cyberrealistic Pony]</a>
<a href="https://huggingface.co/TheImposterImposters/tamePony_v25" target="_blank">[Tame Pony]</a><br>
<a href="https://civitai.com/models/372465/pony-realism" target="_blank">[Pony Realism civitai]</a>
<a href="https://civitai.com/models/443821?modelVersionId=680915" target="_blank">[Cyberrealistic Pony civitai]</a>
<a href="https://civitai.com/models/722045/tame-pony-the-authenticity-machine" target="_blank">[Tame Pony civitai]</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):
            model_choice = gr.Dropdown(
                    ["Pony Realism v22", "Cyber Realistic Pony v65", "Tame Pony v25"],
                    label="Model Choice",
                    value="Pony Realism v22")
            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, 100, 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="Euler a")
                clip_skip = gr.Slider(1, 4, 2, step=1, label="Clip skip")

            with gr.Accordion("Enhancers", open=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.Accordion("Prefix and Suffix Settings", open=True):
                use_positive_prefix = gr.Checkbox(label="Use Positive Prefix", value=True, info=f"Prefix: {DEFAULT_POSITIVE_PREFIX}")
                use_positive_suffix = gr.Checkbox(label="Use Positive Suffix", value=True, info=f"Suffix: {DEFAULT_POSITIVE_SUFFIX}")
                use_negative_prefix = gr.Checkbox(label="Use Negative Prefix", value=True, info=f"Prefix: {DEFAULT_NEGATIVE_PREFIX}")
                use_negative_suffix = gr.Checkbox(label="Use Negative Suffix", value=True, info=f"Suffix: {DEFAULT_NEGATIVE_SUFFIX}")

        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_positive_prompt_used = gr.Textbox(label="Full Positive Prompt Used")
            full_negative_prompt_used = gr.Textbox(label="Full Negative Prompt Used")

    generate_btn.click(
        fn=generate_image,
        inputs=[
            model_choice,
            positive_prompt, negative_prompt, height, width, num_inference_steps,
            guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
            clip_skip, use_medium_enhancer, use_long_enhancer,
            use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
        ],
        outputs=[output_gallery, seed_used, full_positive_prompt_used, full_negative_prompt_used]
    )

demo.launch(debug=True, mcp_server=True)