import os import json import copy import time import requests import random import logging import numpy as np import spaces from typing import Any, Dict, List, Optional, Union import torch from PIL import Image import gradio as gr from diffusers import ( DiffusionPipeline, AutoencoderKL, ZImagePipeline ) from huggingface_hub import ( hf_hub_download, HfFileSystem, ModelCard, snapshot_download) from diffusers.utils import load_image from typing import Iterable from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes colors.orange_red = colors.Color( name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366", c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000", c900="#992900", c950="#802200", ) class OrangeRedTheme(Soft): def __init__( self, *, primary_hue: colors.Color | str = colors.gray, secondary_hue: colors.Color | str = colors.orange_red, # Use the new color neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg, font: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("Outfit"), "Arial", "sans-serif", ), font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", ), ): super().__init__( primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono, ) super().set( background_fill_primary="*primary_50", background_fill_primary_dark="*primary_900", body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", button_primary_text_color="white", button_primary_text_color_hover="white", button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", button_secondary_text_color="black", button_secondary_text_color_hover="white", button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)", button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)", button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)", button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)", slider_color="*secondary_500", slider_color_dark="*secondary_600", block_title_text_weight="600", block_border_width="3px", block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg", button_large_padding="11px", color_accent_soft="*primary_100", block_label_background_fill="*primary_200", ) orange_red_theme = OrangeRedTheme() loras = [ { "image": "https://huggingface.co/Ttio2/Z-Image-Turbo-pencil-sketch/resolve/main/images/z-image_00097_.png", "title": "Turbo Pencil", "repo": "Ttio2/Z-Image-Turbo-pencil-sketch", #0 "weights": "Zimage_pencil_sketch.safetensors", "trigger_word": "pencil sketch" }, { "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/1111111111.png", "title": "AWPortrait Z", "repo": "Shakker-Labs/AWPortrait-Z", #1 "weights": "AWPortrait-Z.safetensors", "trigger_word": "Portrait" }, { "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/cookie-mons.png", "title": "Yarn Art Style", "repo": "linoyts/yarn-art-style", #28 "weights": "yarn-art-style_000001250.safetensors", "trigger_word": "yarn art style" }, { "image": "https://huggingface.co/Quorlen/Z-Image-Turbo-Behind-Reeded-Glass-Lora/resolve/main/images/ComfyUI_00391_.png", "title": "Behind Reeded Glass", "repo": "Quorlen/Z-Image-Turbo-Behind-Reeded-Glass-Lora", #26 "weights": "Z_Image_Turbo_Behind_Reeded_Glass_Lora_TAV2_000002750.safetensors", "trigger_word": "Act1vate!, Behind reeded glass" }, { "image": "https://huggingface.co/ostris/z_image_turbo_childrens_drawings/resolve/main/images/1764433619736__000003000_9.jpg", "title": "Childrens Drawings", "repo": "ostris/z_image_turbo_childrens_drawings", #2 "weights": "z_image_turbo_childrens_drawings.safetensors", "trigger_word": "Children Drawings" }, { "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/xcxc.png", "title": "Tarot Z", "repo": "multimodalart/tarot-z-image-lora", #22 "weights": "tarot-z-image_000001250.safetensors", "trigger_word": "trtcrd" }, { "image": "https://huggingface.co/renderartist/Technically-Color-Z-Image-Turbo/resolve/main/images/ComfyUI_00917_.png", "title": "Technically Color Z", "repo": "renderartist/Technically-Color-Z-Image-Turbo", #3 "weights": "Technically_Color_Z_Image_Turbo_v1_renderartist_2000.safetensors", "trigger_word": "t3chnic4lly" }, { "image": "https://huggingface.co/strangerzonehf/Flux-Ultimate-LoRA-Collection/resolve/main/images/z-image_00147_.png", "title": "Turbo Ghibli", "repo": "Ttio2/Z-Image-Turbo-Ghibli-Style", #19 "weights": "ghibli_zimage_finetune.safetensors", "trigger_word": "Ghibli Style" }, { "image": "https://huggingface.co/tarn59/pixel_art_style_lora_z_image_turbo/resolve/main/images/ComfyUI_00273_.png", "title": "Pixel Art", "repo": "tarn59/pixel_art_style_lora_z_image_turbo", #4 "weights": "pixel_art_style_z_image_turbo.safetensors", "trigger_word": "Pixel art style." }, { "image": "https://huggingface.co/renderartist/Saturday-Morning-Z-Image-Turbo/resolve/main/images/Saturday_Morning_Z_15.png", "title": "Saturday Morning", "repo": "renderartist/Saturday-Morning-Z-Image-Turbo", #5 "weights": "Saturday_Morning_Z_Image_Turbo_v1_renderartist_1250.safetensors", "trigger_word": "saturd4ym0rning" }, { "image": "https://huggingface.co/AIImageStudio/ReversalFilmGravure_z_Image_turbo/resolve/main/images/2025-12-01_173047-z_image_z_image_turbo_bf16-435125750859057-euler_10_hires.png", "title": "ReversalFilmGravure", "repo": "AIImageStudio/ReversalFilmGravure_z_Image_turbo", #6 "weights": "z_image_turbo_ReversalFilmGravure_v1.0.safetensors", "trigger_word": "Reversal Film Gravure, analog film photography" }, { "image": "https://huggingface.co/renderartist/Coloring-Book-Z-Image-Turbo-LoRA/resolve/main/images/CBZ_00274_.png", "title": "Coloring Book Z", "repo": "renderartist/Coloring-Book-Z-Image-Turbo-LoRA", #7 "weights": "Coloring_Book_Z_Image_Turbo_v1_renderartist_2000.safetensors", "trigger_word": "c0l0ringb00k" }, { "image": "https://huggingface.co/damnthatai/1950s_American_Dream/resolve/main/images/ZImage_20251129163459_135x_00001_.jpg", "title": "1950s American Dream", "repo": "damnthatai/1950s_American_Dream", #8 "weights": "5os4m3r1c4n4_z.safetensors", "trigger_word": "5os4m3r1c4n4, 1950s, painting, a painting of" }, { "image": "https://huggingface.co/wcde/Z-Image-Turbo-DeJPEG-Lora/resolve/main/images/01.png", "title": "DeJPEG", "repo": "wcde/Z-Image-Turbo-DeJPEG-Lora", #9 "weights": "dejpeg_v3.safetensors", "trigger_word": "" }, { "image": "https://huggingface.co/suayptalha/Z-Image-Turbo-Realism-LoRA/resolve/main/images/n4aSpqa-YFXYo4dtcIg4W.png", "title": "DeJPEG", "repo": "suayptalha/Z-Image-Turbo-Realism-LoRA", #10 "weights": "pytorch_lora_weights.safetensors", "trigger_word": "Realism" }, { "image": "https://huggingface.co/renderartist/Classic-Painting-Z-Image-Turbo-LoRA/resolve/main/images/Classic_Painting_Z_00247_.png", "title": "Classic Painting Z", "repo": "renderartist/Classic-Painting-Z-Image-Turbo-LoRA", #11 "weights": "Classic_Painting_Z_Image_Turbo_v1_renderartist_1750.safetensors", "trigger_word": "class1cpa1nt" }, { "image": "https://huggingface.co/DK9/3D_MMORPG_style_z-image-turbo_lora/resolve/main/images/10_with_lora.png", "title": "3D MMORPG", "repo": "DK9/3D_MMORPG_style_z-image-turbo_lora", #12 "weights": "lostark_v1.safetensors", "trigger_word": "" }, { "image": "https://huggingface.co/Danrisi/Olympus_UltraReal_ZImage/resolve/main/images/Z-Image_01011_.png", "title": "Olympus UltraReal", "repo": "Danrisi/Olympus_UltraReal_ZImage", #13 "weights": "Olympus.safetensors", "trigger_word": "digital photography, early 2000s compact camera aesthetic, amateur candid shot, digital photography, early 2000s compact camera aesthetic, amateur candid shot, direct flash lighting, hard flash shadow, specular highlights, overexposed highlights" }, { "image": "https://huggingface.co/AiAF/D-ART_Z-Image-Turbo_LoRA/resolve/main/images/example_l3otpwzaz.png", "title": "D ART Z Image", "repo": "AiAF/D-ART_Z-Image-Turbo_LoRA", #14 "weights": "D-ART_Z-Image-Turbo.safetensors", "trigger_word": "D-ART" }, { "image": "https://huggingface.co/AlekseyCalvin/Marionette_Modernism_Z-image-Turbo_LoRA/resolve/main/bluebirdmandoll.webp", "title": "Marionette Modernism", "repo": "AlekseyCalvin/Marionette_Modernism_Z-image-Turbo_LoRA", #15 "weights": "ZImageDadadoll_000003600.safetensors", "trigger_word": "DADADOLL style" }, { "image": "https://huggingface.co/AlekseyCalvin/HistoricColor_Z-image-Turbo-LoRA/resolve/main/HSTZgen2.webp", "title": "Historic Color Z", "repo": "AlekseyCalvin/HistoricColor_Z-image-Turbo-LoRA", #16 "weights": "ZImage1HST_000004000.safetensors", "trigger_word": "HST style" }, { "image": "https://huggingface.co/tarn59/80s_air_brush_style_z_image_turbo/resolve/main/images/ComfyUI_00707_.png", "title": "80s Air Brush", "repo": "tarn59/80s_air_brush_style_z_image_turbo", #17 "weights": "80s_air_brush_style_v2_z_image_turbo.safetensors", "trigger_word": "80s Air Brush style." }, { "image": "https://huggingface.co/CedarC/Z-Image_360/resolve/main/images/1765505225357__000006750_6.jpg", "title": "360panorama", "repo": "CedarC/Z-Image_360", #18 "weights": "Z-Image_360.safetensors", "trigger_word": "360panorama" }, { "image": "https://huggingface.co/jj4real/zimage-igbaddie/resolve/main/images/bc73263da73225d636a1677858435a64.jpg", "title": "Turbo IG baddie", "repo": "jj4real/zimage-igbaddie", #20 "weights": "zimage-igbaddie.safetensors", "trigger_word": "igbaddie" }, { "image": "https://huggingface.co/HAV0X1014/Z-Image-Turbo-KF-Bat-Eared-Fox-LoRA/resolve/main/images/ComfyUI_00132_.png", "title": "KF-Bat-Eared", "repo": "HAV0X1014/Z-Image-Turbo-KF-Bat-Eared-Fox-LoRA", #21 "weights": "z-image-turbo-bat_eared_fox.safetensors", "trigger_word": "bat_eared_fox_kemono_friends" }, { "image": "https://cdn-uploads.huggingface.co/production/uploads/653cd3049107029eb004f968/IHttgddXu6ZBMo7eyy8p6.png", "title": "80s Horror", "repo": "neph1/80s_horror_movies_lora_zit", #23 "weights": "80s_horror_z_80.safetensors", "trigger_word": "80s_horror" }, { "image": "https://huggingface.co/Quorlen/z_image_turbo_Sunbleached_Protograph_Style_Lora/resolve/main/images/ComfyUI_00024_.png", "title": "Sunbleached Protograph", "repo": "Quorlen/z_image_turbo_Sunbleached_Protograph_Style_Lora", #24 "weights": "zimageturbo_Sunbleach_Photograph_Style_Lora_TAV2_000002750.safetensors", "trigger_word": "Act1vate!" }, { "image": "https://huggingface.co/bunnycore/Z-Art-2.1/resolve/main/images/ComfyUI_00069_.png", "title": "Z-Art-2.1", "repo": "bunnycore/Z-Art-2.1", #25 "weights": "Z-Image-Art2.1.safetensors", "trigger_word": "anime art" }, { "image": "https://huggingface.co/cactusfriend/longfurby-z/resolve/main/images/1764658860954__000003000_1.jpg", "title": "Longfurby", "repo": "cactusfriend/longfurby-z", #27 "weights": "longfurbyZ.safetensors", "trigger_word": "" }, { "image": "https://huggingface.co/SkyAsl/Pixel-artist-Z/resolve/main/pixel-art-result.png", "title": "Pixel Art", "repo": "SkyAsl/Pixel-artist-Z", #29 "weights": "adapter_model.safetensors", "trigger_word": "a pixel art character" }, ] dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" base_model = "Tongyi-MAI/Z-Image-Turbo" print(f"Loading {base_model} pipeline...") # Initialize Pipeline pipe = ZImagePipeline.from_pretrained( base_model, torch_dtype=dtype, low_cpu_mem_usage=False, ).to(device) # ======== AoTI compilation + FA3 ======== # As per reference for optimization try: print("Applying AoTI compilation and FA3...") pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"] spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3") print("Optimization applied successfully.") except Exception as e: print(f"Optimization warning: {e}. Continuing with standard pipeline.") MAX_SEED = np.iinfo(np.int32).max class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") def update_selection(evt: gr.SelectData, width, height): selected_lora = loras[evt.index] new_placeholder = f"Type a prompt for {selected_lora['title']}" lora_repo = selected_lora["repo"] updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅" if "aspect" in selected_lora: if selected_lora["aspect"] == "portrait": width = 768 height = 1024 elif selected_lora["aspect"] == "landscape": width = 1024 height = 768 else: width = 1024 height = 1024 return ( gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height, ) @spaces.GPU def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)): # Clean up previous LoRAs in both cases with calculateDuration("Unloading LoRA"): pipe.unload_lora_weights() # Check if a LoRA is selected if selected_index is not None and selected_index < len(loras): selected_lora = loras[selected_index] lora_path = selected_lora["repo"] trigger_word = selected_lora["trigger_word"] # Prepare Prompt with Trigger Word if(trigger_word): if "trigger_position" in selected_lora: if selected_lora["trigger_position"] == "prepend": prompt_mash = f"{trigger_word} {prompt}" else: prompt_mash = f"{prompt} {trigger_word}" else: prompt_mash = f"{trigger_word} {prompt}" else: prompt_mash = prompt # Load LoRA with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"): weight_name = selected_lora.get("weights", None) try: pipe.load_lora_weights( lora_path, weight_name=weight_name, adapter_name="default", low_cpu_mem_usage=True ) # Set adapter scale pipe.set_adapters(["default"], adapter_weights=[lora_scale]) except Exception as e: print(f"Error loading LoRA: {e}") gr.Warning("Failed to load LoRA weights. Generating with base model.") else: # Base Model Case print("No LoRA selected. Running with Base Model.") prompt_mash = prompt with calculateDuration("Randomizing seed"): if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) # Note: Z-Image-Turbo is strictly T2I in this reference implementation. # Img2Img via image_input is disabled/ignored for this pipeline update. with calculateDuration("Generating image"): # For Turbo models, guidance_scale is typically 0.0 forced_guidance = 0.0 # Turbo mode final_image = pipe( prompt=prompt_mash, height=int(height), width=int(width), num_inference_steps=int(steps), guidance_scale=forced_guidance, generator=generator, ).images[0] yield final_image, seed, gr.update(visible=False) def get_huggingface_safetensors(link): split_link = link.split("/") if(len(split_link) == 2): model_card = ModelCard.load(link) base_model = model_card.data.get("base_model") print(base_model) # Relaxed check to allow Z-Image or Flux or others, assuming user knows what they are doing # or specifically check for Z-Image-Turbo if base_model not in ["Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.1-dev"]: # Just a warning instead of error to allow experimentation print("Warning: Base model might not match.") image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None) trigger_word = model_card.data.get("instance_prompt", "") image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None fs = HfFileSystem() try: list_of_files = fs.ls(link, detail=False) for file in list_of_files: if(file.endswith(".safetensors")): safetensors_name = file.split("/")[-1] if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))): image_elements = file.split("/") image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}" except Exception as e: print(e) gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA") raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA") return split_link[1], link, safetensors_name, trigger_word, image_url def check_custom_model(link): if(link.startswith("https://")): if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")): link_split = link.split("huggingface.co/") return get_huggingface_safetensors(link_split[1]) else: return get_huggingface_safetensors(link) def add_custom_lora(custom_lora): global loras if(custom_lora): try: title, repo, path, trigger_word, image = check_custom_model(custom_lora) print(f"Loaded custom LoRA: {repo}") card = f'''
"+trigger_word+" as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}