Upload app.py

updated the app

app.py

@@ -1,3 +1,4 @@
+import random
 import torch
 import gradio as gr
 import spaces
@@ -5,12 +6,15 @@ import spaces
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
-from diffusers import ZImagePipeline, ZImageTransformer2DModel  # Z-Image specific
+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():
@@ -18,7 +22,6 @@ def load_pipe():
     if pipe is not None:
         return pipe
 
-    # Load base components like the official demo
     transformer = ZImageTransformer2DModel.from_pretrained(
         BASE_ID,
         subfolder="transformer",
@@ -31,7 +34,6 @@ def load_pipe():
         torch_dtype=torch.bfloat16,
     ).to("cuda")
 
-    # Now load your custom denoiser weights into the transformer
     ckpt_path = hf_hub_download(CUSTOM_REPO, CUSTOM_FILE)
     state = load_file(ckpt_path)
 
@@ -44,31 +46,41 @@ def load_pipe():
     return pipe
 
 @spaces.GPU
-def generate(prompt, steps=9, height=1024, width=1024, seed=0):
+def generate(prompt, height, width):
     p = load_pipe()
-    g = torch.Generator("cuda").manual_seed(int(seed))
+
+    # 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=int(steps),
-        guidance_scale=0.0,  # turbo-style in the official demo
+        num_inference_steps=FIXED_STEPS,
+        guidance_scale=GUIDANCE,
         generator=g,
     ).images[0]
-    return img
+
+    return img, seed
 
 with gr.Blocks() as demo:
-    gr.Markdown("RAYZIST! A Z-Image Turbo Finetune")
-    prompt = gr.Textbox(label="Prompt", lines=5)
-    steps  = gr.Slider(1, 12, value=8, step=1, label="Steps")
-    width  = gr.Dropdown([512, 768, 1024, 1280], value=1024, label="Width")
-    height = gr.Dropdown([512, 768, 1024, 1280], value=1024, label="Height")
-    seed   = gr.Number(value=0, label="Seed")
-    out    = gr.Image(label="Result")
-
-    btn = gr.Button("GO >")
-    btn.click(generate, [prompt, steps, height, width, seed], out)
+    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()