Spaces:
Sleeping
Sleeping
File size: 2,264 Bytes
c192178 efc6236 c192178 efc6236 c192178 efc6236 c192178 efc6236 c192178 efc6236 c192178 efc6236 c192178 efc6236 c192178 efc6236 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import random
import torch
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from diffusers import ZImagePipeline, ZImageTransformer2DModel
BASE_ID = "Tongyi-MAI/Z-Image-Turbo"
CUSTOM_REPO = "MutantSparrow/Ray"
CUSTOM_FILE = "Z-IMAGE-TURBO/Rayzist.v1.0.safetensors"
FIXED_STEPS = 8
GUIDANCE = 1.0
pipe = None
def load_pipe():
global pipe
if pipe is not None:
return pipe
transformer = ZImageTransformer2DModel.from_pretrained(
BASE_ID,
subfolder="transformer",
torch_dtype=torch.bfloat16,
)
pipe = ZImagePipeline.from_pretrained(
BASE_ID,
transformer=transformer,
torch_dtype=torch.bfloat16,
).to("cuda")
ckpt_path = hf_hub_download(CUSTOM_REPO, CUSTOM_FILE)
state = load_file(ckpt_path)
missing, unexpected = pipe.transformer.load_state_dict(state, strict=False)
print("Loaded custom weights.")
print("Missing keys:", len(missing))
print("Unexpected keys:", len(unexpected))
pipe.set_progress_bar_config(disable=True)
return pipe
@spaces.GPU
def generate(prompt, height, width):
p = load_pipe()
# Random seed every run
seed = random.randint(0, 2**31 - 1)
g = torch.Generator("cuda").manual_seed(seed)
img = p(
prompt=prompt,
height=int(height),
width=int(width),
num_inference_steps=FIXED_STEPS,
guidance_scale=GUIDANCE,
generator=g,
).images[0]
return img, seed
with gr.Blocks() as demo:
gr.Markdown("Ray's Z-Image Turbo finetune: RAYZIST!")
prompt = gr.Textbox(label="Prompt", lines=3)
width = gr.Dropdown([512, 768, 1024, 1280, 1344], value=1024, label="Width")
height = gr.Dropdown([512, 768, 1024, 1280, 1344], value=1024, label="Height")
# Button ABOVE output
btn = gr.Button("GO>")
out = gr.Image(label="Your image")
seed_info = gr.Markdown()
def _run(prompt, height, width):
img, seed = generate(prompt, height, width)
return img, f"Seed: `{seed}`"
btn.click(_run, [prompt, height, width], [out, seed_info])
demo.queue()
demo.launch()
|