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, use_free_unip): """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 = use_free_unip) 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 = use_free_unip 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, use_free_unip, 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, use_free_unip) 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"") 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", ) use_free_unip = gr.Checkbox( label="Use free Uni-P predictor", value=False, ) 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, use_free_unip, ], outputs=[result, result_30_steps, seed], ) if __name__ == "__main__": demo.launch()