app.py

import gradio as gr
import numpy as np
import random
import json
import spaces #[uncomment to use ZeroGPU]
from diffusers import (
    AutoencoderKL,
    StableDiffusionXLPipeline,
    DPMSolverMultistepScheduler
)
from huggingface_hub import login, hf_hub_download
from PIL import Image
# from huggingface_hub import login
from SVDNoiseUnet import NPNet64
import functools
import random
from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
import torch
import torch.nn as nn
from einops import rearrange
from torchvision.utils import make_grid
import time
from pytorch_lightning import seed_everything
from torch import autocast
from contextlib import contextmanager, nullcontext
import accelerate
import torchsde
from SVDNoiseUnet import NPNet128
from tqdm import tqdm, trange
from itertools import islice
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "Lykon/dreamshaper-xl-1-0"  # Replace to the model you would like to use
from sampler import UniPCSampler
from customed_unipc_scheduler import CustomedUniPCMultistepScheduler
from spandrel import ModelLoader

precision_scope = autocast

# 1. Define image conversion functions
def pil_image_to_torch_bgr(img: Image.Image) -> torch.Tensor:
    """Convert a PIL image (RGB) to a torch tensor (BGR, uint8 -> float)."""
    img = np.array(img.convert("RGB"))
    img = img[:, :, ::-1]  # Flip RGB to BGR
    img = img.astype(np.float32) / 255.0  # Normalize to [0, 1]
    img = np.transpose(img, (2, 0, 1))  # HWC to CHW
    return torch.from_numpy(img.copy()).unsqueeze(0) # Add batch dimension

def torch_bgr_to_pil_image(tensor: torch.Tensor) -> Image.Image:
    """Convert a torch tensor (BGR, float) to a PIL image (RGB)."""
    tensor = tensor.squeeze(0).clamp(0, 1) # Remove batch dimension and clamp
    img = tensor.detach().cpu().numpy()
    img = np.transpose(img, (1, 2, 0)) # CHW to HWC
    img = img[:, :, ::-1] # Flip BGR to RGB
    img = (img * 255.0).astype(np.uint8)
    return Image.fromarray(img)



def extract_into_tensor(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def append_zero(x):
    return torch.cat([x, x.new_zeros([1])])

def prepare_sdxl_pipeline_step_parameter( pipe: StableDiffusionXLPipeline
                                         , prompts
                                         , need_cfg
                                         , device
                                         , negative_prompt = None
                                         , W = 1024
                                         , H = 1024): # need to correct the format
        (
            prompt_embeds,
            negative_prompt_embeds,
            pooled_prompt_embeds,
            negative_pooled_prompt_embeds,
        ) = pipe.encode_prompt(
            prompt=prompts,
            negative_prompt=negative_prompt,
            device=device,
            do_classifier_free_guidance=need_cfg,
        )
    # timesteps = pipe.scheduler.timesteps
    
        prompt_embeds = prompt_embeds.to(device)
        add_text_embeds = pooled_prompt_embeds.to(device)
        original_size = (W, H)
        crops_coords_top_left = (0, 0)
        target_size = (W, H)
        text_encoder_projection_dim = None
        add_time_ids = list(original_size + crops_coords_top_left + target_size)
        if pipe.text_encoder_2 is None:
            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
        else:
            text_encoder_projection_dim = pipe.text_encoder_2.config.projection_dim
        passed_add_embed_dim = (
            pipe.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
        )
        expected_add_embed_dim = pipe.unet.add_embedding.linear_1.in_features
        if expected_add_embed_dim != passed_add_embed_dim:
            raise ValueError(
                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
            )
        add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype)
        add_time_ids = add_time_ids.to(device)
        negative_add_time_ids = add_time_ids

        if need_cfg:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
        ret_dict = {
            "text_embeds": add_text_embeds,
            "time_ids": add_time_ids
        }
        return prompt_embeds, ret_dict


# New helper to load a list-of-dicts preference JSON
# JSON schema: [ { 'human_preference': [int], 'prompt': str, 'file_path': [str] }, ... ]
def load_preference_json(json_path: str) -> list[dict]:
    """Load records from a JSON file formatted as a list of preference dicts."""
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data

# New helper to extract just the prompts from the preference JSON
# Returns a flat list of all 'prompt' values

def extract_prompts_from_pref_json(json_path: str) -> list[str]:
    """Load a JSON of preference records and return only the prompts."""
    records = load_preference_json(json_path)
    return [rec['prompt'] for rec in records]

# Example usage:
# prompts = extract_prompts_from_pref_json("path/to/preference.json")
# print(prompts)

def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu',need_append_zero = True):
    """Constructs the noise schedule of Karras et al. (2022)."""
    ramp = torch.linspace(0, 1, n)
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    return append_zero(sigmas).to(device) if need_append_zero else sigmas.to(device)

def extract_into_tensor(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))

def append_zero(x):
    return torch.cat([x, x.new_zeros([1])])

def append_dims(x, target_dims):
    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
    dims_to_append = target_dims - x.ndim
    if dims_to_append < 0:
        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
    return x[(...,) + (None,) * dims_to_append]


def chunk(it, size):
    it = iter(it)
    return iter(lambda: tuple(islice(it, size)), ())

def convert_caption_json_to_str(json):
    caption = json["caption"]
    return caption


DTYPE = torch.float16  # torch.float16 works as well, but pictures seem to be a bit worse
device = "cuda" 
cyberreal_repo = "cyberdelia/CyberRealisticXL"
cyberreal_filename = "CyberRealisticXLPlay_V7.0_FP16.safetensors"
cyberreal_path = hf_hub_download(
        repo_id=cyberreal_repo,
        filename=cyberreal_filename,
        cache_dir="."
)

pipe = StableDiffusionXLPipeline.from_single_file(
        cyberreal_path,
        torch_dtype=DTYPE,
)

up_repo = "uwg/upscaler"
up_filename = "ESRGAN/4x_NMKD-Siax_200k.pth"
up_path = hf_hub_download(
        repo_id=up_repo,
        filename=up_filename,
        cache_dir="."
)
upscaler = ModelLoader().load_from_file(up_path)
upscaler.to(device).eval() 



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

accelerator = accelerate.Accelerator()

def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps):
    """Helper function to generate image with specific number of steps"""
    scheduler = CustomedUniPCMultistepScheduler.from_config(pipe.scheduler.config
                                                            , solver_order = 2 if num_inference_steps==8 else 1
                                                            ,denoise_to_zero = False
                                                            , use_afs = True
                                                            , use_free_predictor = False)
    start_free_at_step = 4
    pipe.scheduler = scheduler
    pipe.to('cuda')
    with torch.no_grad():
        with precision_scope("cuda"):
            prompts =  [prompt]
            
            latents = torch.randn(
                (1, pipe.unet.config.in_channels, height // 8, width // 8),
                device=device,
            )
            latents = latents * pipe.scheduler.init_noise_sigma
                    
            pipe.scheduler.set_timesteps(num_inference_steps)
            idx = 0
            register_free_upblock2d(pipe, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
            register_free_crossattn_upblock2d(pipe, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
            for t in tqdm(pipe.scheduler.timesteps):
                        # Still not enough. I will tell you, what is the best implementation.  Although not via the following code.
                        
                        # if idx == len(pipe.scheduler.timesteps) - 1: 
                        #     break
                if idx == start_free_at_step:#(6 if num_inference_steps == 8 else 4):
                    register_free_upblock2d(pipe, b1=1.2, b2=1.2, s1=0.9, s2=0.9)
                    register_free_crossattn_upblock2d(pipe, b1=1.2, b2=1.2, s1=0.9, s2=0.9)
                latent_model_input  = torch.cat([latents] * 2)
                        
                latent_model_input  = pipe.scheduler.scale_model_input(latent_model_input , timestep=t)
                negative_prompts = 'lowres, bad anatomy, bad hands, watermark'
                negative_prompts = 1 * [negative_prompts]
                use_afs = True
                use_free_predictor = False
                prompt_embeds, cond_kwargs = prepare_sdxl_pipeline_step_parameter(pipe
                                                                                      , prompts
                                                                                      , need_cfg=True
                                                                                      , device=pipe.device
                                                                                      , negative_prompt=negative_prompts
                                                                                      , W=width
                                                                                      , H=height)
                if idx == 0 and use_afs:
                    noise_pred = latent_model_input * 0.98
                elif idx == len(pipe.scheduler.timesteps) - 1 and use_free_predictor:
                    noise_pred = None
                else:
                    noise_pred  = pipe.unet(latent_model_input 
                                        , t
                                        , encoder_hidden_states=prompt_embeds.to(device=latents.device, dtype=latents.dtype)
                                        , added_cond_kwargs=cond_kwargs).sample
                if noise_pred is not None:
                    uncond, cond = noise_pred.chunk(2)
                    noise_pred  = uncond + (cond - uncond) * guidance_scale
                latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample
                idx += 1
                        
            x_samples_ddim = pipe.vae.decode(latents / pipe.vae.config.scaling_factor).sample
            x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
            if True:
                for x_sample in x_samples_ddim:
                                # x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                    x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                    img = Image.fromarray(x_sample.astype(np.uint8))#.save( os.path.join(sample_path, f"{base_count:05}.png"))
                    input_image_tensor = pil_image_to_torch_bgr(img).to(device)
                    output_tensor = upscaler(input_image_tensor)
                    output_image_pil = torch_bgr_to_pil_image(output_tensor)
    return output_image_pil

@spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    resolution,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Parse resolution string into width and height
    width, height = map(int, resolution.split('x'))
    
    # Generate image with selected steps
    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps)
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config
                                                             , final_sigmas_type="sigma_min"
                                                             , algorithm_type="sde-dpmsolver++"
                                                             , use_karras_sigmas=True)
    # Generate image with 50 steps for high quality
    negative_prompts = 'lowres, bad anatomy, bad hands, watermark'
    negative_prompts = 1 * [negative_prompts]
    image_50_steps = pipe(prompt=[prompt]
                            ,negative_prompt=negative_prompts
                            ,num_inference_steps=30
                            ,guidance_scale=4.0
                            ,height=height
                            ,width=width).images
    for x_sample in image_50_steps:
        input_image_tensor = pil_image_to_torch_bgr(x_sample).to(device)
        output_tensor = upscaler(input_image_tensor)
        img_4k_org = torch_bgr_to_pil_image(output_tensor)
    return image_quick, img_4k_org, seed


examples = [
    "ultra-realistic 8k RAW portrait of a serious Black man in 1920s Harlem, standing on a bustling vintage city street, wearing a textured vintage wool suit, striped dress shirt, bold colorful tie, and a brown felt fedora, cinematic lighting with soft shadows on his deeply expressive face, timeless and melancholic mood, blurred storefronts and pedestrians in background, analog film grain, slightly desaturated color palette, medium format lens capturing fine skin texture, worn fabric, and atmospheric detail, Harlem Renaissance style, captured in natural light, shallow depth of field",
    "An ultra-realistic 8k HDR editorial photograph of a soft-featured young woman with auburn hair tucked under a linen bonnet, pale freckled skin and downcast eyes filled with quiet resilience, dressed in a modest 1875 working-class Victorian dress with worn shawl, standing near a bustling street market in London, surrounded by wooden carts, hanging meats, and soot-stained brick buildings, soft overcast light and rising chimney smoke blending into a hazy amber atmosphere, cinematic lens depth with visible film grain and rich Kodak Portra-style color grading, historical fashion editorial with immersive composition and a contemplative, narrative mood",
    "A weathered Victorian house surrounded by lush autumn foliage and overgrown garden paths, its deep teal-painted wood faded and peeling, orange leaves scattering across the stone steps and tangled in the railings of the ornate wooden porch, delicate orange wildflowers growing from cracks in the stairs, arched twin doors with stained glass glowing faintly from within, warm golden light filtering through dusted windows, a few butterflies fluttering through the crisp autumn air, the scene bathed in soft daylight with painterly shadows, magical realism meets gothic nostalgia, cinematic composition with high detail and storybook charm, photorealistic yet slightly stylized, peaceful and enchanted with a hint of mystery",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks() as demo:
    gr.HTML(f"<style>{css}</style>")
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Hyperparameters are all you need")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )

            run_button = gr.Button("Run", scale=0, variant="primary")

        with gr.Row():
            with gr.Column():
                gr.Markdown("### Our fast inference Result using afs to get 1 free steps")
                result = gr.Image(label="Quick Result", show_label=False)
            with gr.Column():
                gr.Markdown("### official 30 steps result")
                result_30_steps = gr.Image(label="30 Steps Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=False,
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            resolution = gr.Dropdown(
                choices=[
                    "1024x1024",
                    "1216x832",
                    "832x1216"
                ],
                value="832x1216",
                label="Resolution",
            )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=5.0,
                    step=0.1,
                    value=5.0,  # Replace with defaults that work for your model
                )

                num_inference_steps = gr.Dropdown(
                    choices=[6, 7, 8],
                    value=8,
                    label="Number of inference steps",
                )

        gr.Examples(examples=examples, inputs=[prompt])
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            resolution,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, result_20_steps, seed],
    )

if __name__ == "__main__":
    demo.launch()