app.py

#!/usr/bin/env python
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
import spaces

import os
import random
import uuid
import gradio as gr
import numpy as np
from PIL import Image
import torch
from diffusers import AutoencoderKL, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
#from diffusers import AutoencoderKL
from typing import Tuple
#from transformers import AutoTokenizer, AutoModelForCausalLM
import paramiko
import gc
#os.system("chmod +x ./cusparselt.sh")
#os.system("./cusparselt.sh")
#os.system("chmod +x ./cudnn.sh")
#os.system("./cudnn.sh")

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
torch.backends.cuda.preferred_blas_library="cublas"
torch.backends.cuda.preferred_linalg_library="cusolver"

torch.set_float32_matmul_precision("highest")

FTP_HOST = "1ink.us"
FTP_USER = "ford442"
FTP_PASS = "GoogleBez12!"
FTP_DIR = "1ink.us/stable_diff/"  # Remote directory on FTP server

css = '''
.gradio-container{max-width: 570px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''

DESCRIPTIONXX = """
    ## REALVISXL V5.0 BF16 ⚡⚡⚡⚡
"""

examples = [

    "Many apples splashed with drops of water within a fancy bowl 4k, hdr  --v 6.0 --style raw",
    "A profile photo of a dog, brown background, shot on Leica M6 --ar 128:85 --v 6.0 --style raw",
]

MODEL_OPTIONS = {
    "REALVISXL V5.0 BF16": "ford442/RealVisXL_V5.0_BF16",
}

MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = 0
BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))

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

style_list = [
    {
        "name": "3840 x 2160",
        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "2560 x 1440",
        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "HD+",
        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "Style Zero",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
DEFAULT_STYLE_NAME = "Style Zero"
STYLE_NAMES = list(styles.keys())

def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    if style_name in styles:
        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    else:
        p, n = styles[DEFAULT_STYLE_NAME]
    if not negative:
        negative = ""
    return p.replace("{prompt}", positive), n + negative

def load_and_prepare_model(model_id):
    model_dtypes = {
        "ford442/RealVisXL_V5.0_BF16": torch.bfloat16,
    }
    dtype = model_dtypes.get(model_id, torch.bfloat16)  # Default to float32 if not found
    vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16,safety_checker=None)
    pipe = StableDiffusionXLPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        add_watermarker=False,
        use_safetensors=True,
        vae=vae,
    ).to(torch.bfloat16).to('cuda')
    
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
    return pipe

# Preload and compile both models
models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}

MAX_SEED = np.iinfo(np.int32).max

def upload_to_ftp(filename):
    try:
        transport = paramiko.Transport((FTP_HOST, 22))
        destination_path=FTP_DIR+filename
        transport.connect(username = FTP_USER, password = FTP_PASS)
        sftp = paramiko.SFTPClient.from_transport(transport)
        sftp.put(filename, destination_path)
        sftp.close()
        transport.close()
        print(f"Uploaded {filename} to FTP server")
    except Exception as e:
        print(f"FTP upload error: {e}")
        
def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name,optimize=False,compress_level=0)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

GPU_DURATION_OPTIONS = {
    "Short (25s)": 25,
    "Short (45s)": 45,
    "Medium (60s)": 60,
    "Medium (80s)": 80,
    "Long (100s)": 100,
    "Long (120s)": 120,
    "Long (140s)": 140,
}

def set_gpu_duration(duration_choice):
    os.environ["GPU_DURATION"] = str(GPU_DURATION_OPTIONS[duration_choice])

@gpu_with_duration(duration=45)
def generate_45s(*args, **kwargs):
    return generate(*args, **kwargs)  # Call the common generate function

@gpu_with_duration(duration=60)
def generate_60s(*args, **kwargs):
    return generate(*args, **kwargs)

@gpu_with_duration(duration=80)
def generate_80s(*args, **kwargs):
    return generate(*args, **kwargs)

@gpu_with_duration(duration=100)
def generate_100s(*args, **kwargs):
    return generate(*args, **kwargs)

def generate(
    model_choice: str,
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style_selection: str = DEFAULT_STYLE_NAME,
    seed: int = 1,
    width: int = 768,
    height: int = 768,
    guidance_scale: float = 3.0,
    num_inference_steps: int = 250,
    randomize_seed: bool = False,
    use_resolution_binning: bool = True, 
    num_images: int = 1,  
    progress=gr.Progress(track_tqdm=True),
):
    global models
    pipe = models[model_choice]
    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator(device='cuda').manual_seed(seed)

    prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)

    options = {
        "prompt": [prompt] * num_images,
        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "generator": generator,
        "output_type": "pil",
    }

    if use_resolution_binning:
        options["use_resolution_binning"] = True

    images = []
    #with torch.no_grad():
    for i in range(0, num_images, BATCH_SIZE):
        batch_options = options.copy()
        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
        if "negative_prompt" in batch_options:
            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
        images.extend(pipe(**batch_options).images)
    sd_image_path = f"rv50_{seed}.png"
    images[0].save(sd_image_path,optimize=False,compress_level=0)
    upload_to_ftp(sd_image_path)    
    image_paths = [save_image(img) for img in images]
    torch.cuda.empty_cache()
    gc.collect()
    return image_paths, seed

def generate_cpu(
    model_choice: str,
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style_selection: str = DEFAULT_STYLE_NAME,
    seed: int = 1,
    width: int = 768,
    height: int = 768,
    guidance_scale: float = 3,
    num_inference_steps: int = 250,
    randomize_seed: bool = False,
    use_resolution_binning: bool = True, 
    num_images: int = 1,  
    progress=gr.Progress(track_tqdm=True),
):
    global models
    pipe = models[model_choice]
    pipe.to("cpu")

    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator(device='cpu').manual_seed(seed) # cpu function seed

    prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)

    options = {
        "prompt": [prompt] * num_images,
        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "generator": generator,
        "output_type": "pil",
    }

    if use_resolution_binning:
        options["use_resolution_binning"] = True

    images = []
    for i in range(0, num_images, BATCH_SIZE):
        batch_options = options.copy()
        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
        if "negative_prompt" in batch_options:
            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
        images.extend(pipe(**batch_options).images)

    image_paths = [save_image(img) for img in images]
    return image_paths, seed

def set_gpu_duration(duration_choice):
    global global_gpu_duration
    global_gpu_duration = GPU_DURATION_OPTIONS[duration_choice]

def load_predefined_images1():
    predefined_images1 = [
        "assets/7.png",
        "assets/8.png",
        "assets/9.png",
        "assets/1.png",
        "assets/2.png",
        "assets/3.png",
        "assets/4.png",
        "assets/5.png",
        "assets/6.png",
    ]
    return predefined_images1


# def load_predefined_images():
#     predefined_images = [
#         "assets2/11.png",
#         "assets2/22.png",
#         "assets2/33.png",
#         "assets2/44.png",
#         "assets2/55.png",
#         "assets2/66.png",
#         "assets2/77.png",
#         "assets2/88.png",
#         "assets2/99.png",
#     ]
#     return predefined_image


with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown(DESCRIPTIONXX)
    with gr.Row():
        prompt = gr.Text(
            label="Prompt",
            show_label=False,
            max_lines=1,
            placeholder="Enter your prompt",
            container=False,
        )
        run_button_45 = gr.Button("Run_45", scale=0)
        cpu_run_button = gr.Button("CPU Run", scale=0)
        run_button_60 = gr.Button("Run_60", scale=0)
        run_button_80 = gr.Button("Run_80", scale=0)
        run_button_100 = gr.Button("Run_100", scale=0)
    result = gr.Gallery(label="Result", columns=1, show_label=False) 

    with gr.Row():
        model_choice = gr.Dropdown(
            label="Model Selection🔻",
            choices=list(MODEL_OPTIONS.keys()),
            value="REALVISXL V5.0 BF16"
        )

    with gr.Accordion("Advanced options", open=False, visible=True):
        gpu_duration = gr.Dropdown(
            label="GPU Duration",
            choices=list(GPU_DURATION_OPTIONS.keys()),
            value="Medium (60s)"  # Default value
        )
        
        gpu_duration.change(fn=set_gpu_duration, inputs=gpu_duration, outputs=[])

        style_selection = gr.Radio(
            show_label=True,
            container=True,
            interactive=True,
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            label="Quality Style",
        )
        num_images = gr.Slider(
            label="Number of Images",
            minimum=1,
            maximum=5,
            step=1,
            value=1,
        )
        with gr.Row():
            with gr.Column(scale=1):
                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
                negative_prompt = gr.Text(
                    label="Negative prompt",
                    max_lines=5,
                    lines=4,
                    placeholder="Enter a negative prompt",
                    value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
                    visible=True,
                )
        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row():
            width = gr.Slider(
                label="Width",
                minimum=448,
                maximum=MAX_IMAGE_SIZE,
                step=64,
                value=768,
            )
            height = gr.Slider(
                label="Height",
                minimum=448,
                maximum=MAX_IMAGE_SIZE,
                step=64,
                value=768,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.1,
                maximum=6,
                step=0.1,
                value=3.0,
            )
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=10,
                maximum=1000,
                step=10,
                value=250,
            )

    gr.Examples(
        examples=examples,
        inputs=prompt,
        cache_examples=False
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )
    
    gr.on(
        triggers=[
            run_button_100.click,
        ],
      #  api_name="generate",  # Add this line
        fn=generate_100s,
        inputs=[
            model_choice,
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            randomize_seed, 
            num_images, 
        ],
        outputs=[result, seed],
    )
    
    gr.on(
        triggers=[
            run_button_80.click,
        ],
      #  api_name="generate",  # Add this line
        fn=generate_80s,
        inputs=[
            model_choice,
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            randomize_seed, 
            num_images, 
        ],
        outputs=[result, seed],
    )
    
    gr.on(
        triggers=[
            run_button_60.click,
        ],
      #  api_name="generate",  # Add this line
        fn=generate_60s,
        inputs=[
            model_choice,
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            randomize_seed, 
            num_images, 
        ],
        outputs=[result, seed],
    )

    gr.on(
        triggers=[
            run_button_45.click,
        ],
      #  api_name="generate",  # Add this line
        fn=generate_45s,
        inputs=[
            model_choice,
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            randomize_seed, 
            num_images, 
        ],
        outputs=[result, seed],
    )

    gr.Markdown("### REALVISXL V5.0")
    predefined_gallery = gr.Gallery(label="REALVISXL V5.0", columns=3, show_label=False, value=load_predefined_images1())

    #gr.Markdown("### LIGHTNING V5.0")
    #predefined_gallery = gr.Gallery(label="LIGHTNING V5.0", columns=3, show_label=False, value=load_predefined_images())

    gr.Markdown(
    """
    <div style="text-align: justify;">
    ⚡Models used in the playground <a href="https://huggingface.co/SG161222/RealVisXL_V5.0">[REALVISXL V5.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V5.0_Lightning">[REALVISXL V5.0 LIGHTNING]</a> for image generation. Stable Diffusion XL piped (SDXL) model HF. This is the demo space for generating images using the Stable Diffusion XL models, with multiple different variants available.
    </div>
    """)

    gr.Markdown(
    """
    <div style="text-align: justify;">
    ⚡This is the demo space for generating images using Stable Diffusion XL with quality styles, different models, and types. Try the sample prompts to generate higher quality images. Try the sample prompts for generating higher quality images. 
    <a href='https://huggingface.co/spaces/prithivMLmods/Top-Prompt-Collection' target='_blank'>Try prompts</a>.
    </div>
    """)

    gr.Markdown(
    """
    <div style="text-align: justify;">
    ⚠️ Users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.
    </div>
    """) 

def text_generation(input_text, seed):
    full_prompt = "Text Generator Application by ecarbo"
    return full_prompt
    
title = "Text Generator Demo GPT-Neo"
description = "Text Generator Application by ecarbo"

if __name__ == "__main__":
    demo_interface = demo.queue(max_size=50)  # Remove .launch() here

    text_gen_interface = gr.Interface(
        fn=text_generation,
        inputs=[
            gr.Textbox(lines=1, label="Expand the following prompt to be more detailed and descriptive for image generation: "),
            gr.Number(value=10, label="Enter seed number")
        ],
        outputs=gr.Textbox(label="Text Generated"),
        title=title,
        description=description,
        theme="huggingface"
    )

    combined_interface = gr.TabbedInterface([demo_interface, text_gen_interface], ["Image Generation", "Text Generation"])
    combined_interface.launch(show_api=False)