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
from diffusers import EulerAncestralDiscreteScheduler
from typing import Tuple
import paramiko
import datetime
# import cyper
from image_gen_aux import UpscaleWithModel

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")

DESCRIPTIONXX = """
    ## ⚡⚡⚡⚡ REALVISXL V5.0 BF16 (Tester L) ⚡⚡⚡⚡
"""

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",
]

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())

HF_TOKEN = os.getenv("HF_TOKEN")
FTP_HOST = os.getenv("FTP_HOST")
FTP_USER = os.getenv("FTP_USER")
FTP_PASS = os.getenv("FTP_PASS")
FTP_DIR = os.getenv("FTP_DIR")

# os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')

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

neg_prompt_2 = " 'non-photorealistic':1.5, 'unrealistic skin','unattractive face':1.3, 'low quality':1.1, ('dull color scheme', 'dull colors', 'digital noise':1.2),'amateurish', 'poorly drawn face':1.3, 'poorly drawn', 'distorted face', 'low resolution', 'simplistic' "

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

upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cpu"))

def load_and_prepare_model():
    vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #.to(device=device, dtype=torch.bfloat16)
    #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler')
    pipe = StableDiffusionXLPipeline.from_pretrained(
        #'ford442/RealVisXL_V5.0_BF16',
        #'ford442/RealVisXL_V5.0_FP64',
        'SG161222/RealVisXL_V5.0',
        trust_remote_code=True,
        add_watermarker=False,
        vae=vaeXL,
        low_cpu_mem_usage = False,
        token=HF_TOKEN,
    )
    #pipe.vae = vaeXL #.to(torch.bfloat16)
    #pipe.scheduler = sched
    #pipe.vae.do_resize=False
    #pipe.vae.vae_scale_factor=8
    #pipe.to(device=device, dtype=torch.bfloat16)
    pipe.to(device)
    #pipe.to(torch.bfloat16)
    #pipe.vae.set_default_attn_processor()
    print(f'init noise scale: {pipe.scheduler.init_noise_sigma}')
    pipe.watermark=None
    pipe.safety_checker=None  
    return pipe
    
pipe = load_and_prepare_model()

l_dtype=torch.float32

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 uploadNote(prompt,num_inference_steps,guidance_scale,timestamp):
    filename= f'rv_C_{timestamp}.txt'
    with open(filename, "w") as f:
        f.write(f"Realvis 5.0 (Tester L) \n")
        f.write(f"Date/time: {timestamp} \n")
        f.write(f"Prompt: {prompt} \n")
        f.write(f"Steps: {num_inference_steps} \n")
        f.write(f"Guidance Scale: {guidance_scale} \n")
    upload_to_ftp(filename) 

@spaces.GPU(duration=60)
def generate(
    segment: int = 1,
    prompt: str = "",
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style_selection: str = "",
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3.8,
    num_inference_steps: int = 200,
    seed: int = 424242,
    use_resolution_binning: bool = True, 
    progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if segment==1:
        seed = random.randint(0, MAX_SEED)
        generator = torch.Generator(device='cuda').manual_seed(seed)
        prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
        prompt=prompt, negative_prompt=negative_prompt, device=device, num_images_per_prompt=1,
        do_classifier_free_guidance=True, prompt_embeds=None, negative_prompt_embeds=None,
        pooled_prompt_embeds=None, negative_pooled_prompt_embeds=None, lora_scale=None, clip_skip=None,
        )
        pipe.scheduler.set_timesteps(num_inference_steps, device=torch.device('cuda'))
        timesteps = pipe.scheduler.timesteps
        all_timesteps_cpu = timesteps.cpu()
        # test with 2 segments
        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
        segment_timesteps = torch.from_numpy(timesteps_split_np[0]).to("cuda")
        # test with 1 segment
        #segment_timesteps = timesteps
        num_channels_latents = pipe.unet.config.in_channels
        latents = pipe.prepare_latents(
            batch_size=1, num_channels_latents=pipe.unet.config.in_channels, height=height, width=width,
            dtype=l_dtype, device=device, generator=generator, latents=None,
        )
        text_encoder_projection_dim = pipe.text_encoder_2.config.projection_dim
        original_size = (width, height)
        target_size = (width, height)
        crops_coords_top_left = (0, 0)
        add_time_ids = pipe._get_add_time_ids(
            original_size, crops_coords_top_left, target_size, dtype=l_dtype, # Use bfloat16
            text_encoder_projection_dim=text_encoder_projection_dim
        ).to("cuda")
        unet_prompt_embeds = prompt_embeds
        unet_added_text_embeds = pooled_prompt_embeds
        loop_add_time_ids = add_time_ids
        unet_prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        unet_added_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds])
        loop_add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
        added_cond_kwargs = {"text_embeds": unet_added_text_embeds, "time_ids": loop_add_time_ids}
        current_latents = latents
    else:
        state_file = f"rv_L_{segment-1}_{seed}.pt"
        state = torch.load(state_file, weights_only=False)
        # # TEST
        #seed = state["seed"]
        generator = torch.Generator(device='cuda') #.manual_seed(seed)
        generator_state = state["generator_state"]
        generator.set_state(generator_state)
        latents = state["intermediate_latents"].to("cuda",dtype=l_dtype) #, dtype=torch.bfloat16)
        guidance_scale = state["guidance_scale"]
        all_timesteps_cpu = state["all_timesteps"]
        height = state["height"]
        width = state["width"]
        pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
        segment_timesteps = torch.from_numpy(timesteps_split_np[segment - 1]).to("cuda")
        prompt_embeds = state["prompt_embeds"].to("cuda", dtype=l_dtype)
        negative_prompt_embeds = state["negative_prompt_embeds"].to("cuda", dtype=l_dtype)
        pooled_prompt_embeds = state["pooled_prompt_embeds"].to("cuda", dtype=l_dtype)
        negative_pooled_prompt_embeds = state["negative_pooled_prompt_embeds"].to("cuda", dtype=l_dtype)
        unet_prompt_embeds = prompt_embeds
        unet_added_text_embeds = pooled_prompt_embeds
        unet_prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        unet_added_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds])
        add_time_ids = state["add_time_ids"].to("cuda", dtype=l_dtype)
        loop_add_time_ids = add_time_ids
        loop_add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
        added_cond_kwargs = {"text_embeds": unet_added_text_embeds, "time_ids": loop_add_time_ids}
        current_latents = latents
    for i, t in enumerate(pipe.progress_bar(segment_timesteps)):
        latent_model_input = torch.cat([current_latents] * 2)
        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
        with torch.no_grad():
          noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=unet_prompt_embeds,added_cond_kwargs=added_cond_kwargs, return_dict=False)[0]
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
        current_latents = pipe.scheduler.step(noise_pred, t, current_latents, return_dict=False)[0]
    intermediate_latents_cpu = current_latents.detach().cpu()
    if segment==8:
        final_latents = current_latents
        final_latents = final_latents / pipe.vae.config.scaling_factor
        #with torch.no_grad():
        image = pipe.vae.decode(final_latents, return_dict=False)[0]
        image = pipe.image_processor.postprocess(image.detach(), output_type="pil")[0]
        output_image_file = f"rv_L_{seed}.png"
        upscaler.to(torch.device('cuda'))
        with torch.no_grad():
            upscale2 = upscaler(image, tiling=True, tile_width=256, tile_height=256)
        #timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        print('-- got upscaled image --')
        downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
        upscale_path = f"rv_L_{seed}.png"
        downscale2.save(upscale_path,optimize=False,compress_level=0)
        upload_to_ftp(upscale_path)
        #timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        #upload_to_ftp(filename) 
        #uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
        unique_name = str(uuid.uuid4()) + ".png"  
        os.symlink(output_image_file, unique_name)  
        return [unique_name], seed
    else:
        original_prompt_embeds_cpu = prompt_embeds.cpu()
        original_negative_prompt_embeds_cpu = negative_prompt_embeds.cpu()
        original_pooled_prompt_embeds_cpu = pooled_prompt_embeds.cpu()
        original_negative_pooled_prompt_embeds_cpu = negative_pooled_prompt_embeds.cpu()
        original_add_time_ids_cpu = add_time_ids.cpu()
        original_generator_state_cpu = generator.get_state().cpu()
        timesteps = pipe.scheduler.timesteps
        all_timesteps_cpu = timesteps.cpu() # Move to CPU
        state = {
        "intermediate_latents": intermediate_latents_cpu,
        "generator_state": original_generator_state_cpu,
        "all_timesteps": all_timesteps_cpu, # Save full list generated by scheduler
        "prompt_embeds": original_prompt_embeds_cpu, # Save ORIGINAL embeds
        "negative_prompt_embeds": original_negative_prompt_embeds_cpu,
        "pooled_prompt_embeds": original_pooled_prompt_embeds_cpu,
        "negative_pooled_prompt_embeds": original_negative_pooled_prompt_embeds_cpu,
        "add_time_ids": original_add_time_ids_cpu, # Save ORIGINAL time IDs
        "guidance_scale": guidance_scale,
        "seed": seed,
        "prompt": prompt, # Save originals for reference/verification
        "negative_prompt": negative_prompt,
        "height": height, # Save dimensions used
        "width": width
        }
        state_file = f"rv_L_{segment}_{seed}.pt"
        torch.save(state, state_file)
        return None, seed

def update_ranges(total_steps):
    """Calculates and updates the ranges for the 8 slave sliders."""
    step_size = total_steps // 8  # Calculate the size of each segment
    ranges = []
    for i in range(8):
        lower_bound = i * step_size
        ranges.append([lower_bound])  # Add the range to the list
    return ranges

with gr.Blocks(theme=gr.themes.Origin()) 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_1 = gr.Button("Run Segment 1", scale=0)
        run_button_2 = gr.Button("Run Segment 2", scale=0)
        run_button_3 = gr.Button("Run Segment 3", scale=0)
        run_button_4 = gr.Button("Run Segment 4", scale=0)
        run_button_5 = gr.Button("Run Segment 5", scale=0)
        run_button_6 = gr.Button("Run Segment 6", scale=0)
        run_button_7 = gr.Button("Run Segment 7", scale=0)
        run_button_8 = gr.Button("Run Segment 8", scale=0)
        result = gr.Gallery(label="Result", columns=1, show_label=False)
        seed = gr.Number(value=1, label="Seed")
    with gr.Row():
        style_selection = gr.Radio(
            show_label=True,
            container=True,
            interactive=True,
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            label="Quality Style",
        )
        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),'not photorealistic':1.5, 'poorly drawn', 'bad anatomy', 'wrong anatomy', 'extra limb', 'missing limb', 'floating limbs', 'poorly drawn hands', 'poorly drawn feet', 'poorly drawn face':1.3, 'out of frame', 'extra limbs', 'bad anatomy', 'bad art', 'beginner',  'distorted face','amateur'",
                    visible=True,
                )
        with gr.Row():
            width = gr.Slider(
                label="Width",
                minimum=448,
                maximum=4096,
                step=64,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=448,
                maximum=4096,
                step=64,
                value=1024,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.1,
                maximum=30,
                step=0.05,
                value=3.8,
            )
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=10,
                maximum=2000,
                step=10,
                value=200,
            )
        range_sliders = []
        for i in range(8):
            slider = gr.Slider(
                minimum=1,
                maximum=250,
                value=[i * (num_inference_steps.value // 8)],
                step=1,
                label=f"Range {i + 1}",
            )
            range_sliders.append(slider)
        num_inference_steps.change(
        update_ranges,
        inputs=num_inference_steps,
        outputs=range_sliders,
        )
        
        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_1.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=1),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_2.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=2),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_3.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=3),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_4.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=4),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_5.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=5),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_6.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=6),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_7.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=7),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    gr.on(
        triggers=[
            run_button_8.click,
        ],
        fn=generate,
        inputs=[
            gr.Number(value=8),
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            seed,
        ],
        outputs=[result, seed],
    )
    
    gr.Markdown("### REALVISXL V5.0 Default Mode")

    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,
    )

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