22

app.py

@@ -1,7 +1,178 @@
 import gradio as gr
+import numpy as np
+import random
+import spaces
+from diffusers import DiffusionPipeline
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+MODEL_REPO = "Tongyi-MAI/Z-Image-Turbo"
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1280
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch(mcp_server=True)
+# Load the pipeline once at startup
+print("Loading Z-Image-Turbo pipeline...")
+pipe = DiffusionPipeline.from_pretrained(
+    MODEL_REPO,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=False,
+)
+pipe.to("cuda")
+
+# ======== AoTI compilation + FA3 ========
+pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+
+@spaces.GPU
+def inference(
+    input_image,
+    prompt,
+    seed=42,
+    randomize_seed=True,
+    width=1024,
+    height=1024,
+    guidance_scale=5.0,
+    num_inference_steps=8,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if input_image is None:
+        print("Error: input_image is empty.")
+        return None
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+
+    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3)
+    pipe.scheduler = scheduler
+
+    ### load input_image
+    
+
+    # input_image = input_image.convert("RGB").resize((1024, 1024))
+
+    # input_tensor = pil_image_to_torch_bgr(input_image).to(pipe.vae.device).to(pipe.vae.dtype)
+    input_tensor = pipe.image_processor.preprocess(input_image, height=height, width=width)
+    # input_tensor = image.to(dtype=torch.float32)
+    input_tensor = input_tensor.to(device=pipe.vae.device, dtype=pipe.vae.dtype)
+
+    # input_latents = _encode_vae_image(pipe=pipe, image=input_tensor, generator=generator)
+    # input_latents = pipe.vae.encode(input_tensor).latent_dist.sample()
+    # input_latents = input_latents * pipe.vae.config.scaling_factor
+    
+    image = pipe(
+        # latents=input_latents,
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        # width=width,
+        # height=height,
+        generator=generator,
+    ).images[0]
+
+    return image, seed
+
+
+def read_file(path: str) -> str:
+    with open(path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    return content
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 640px;
+}
+"""
+
+examples = [
+    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    "An astronaut riding a green horse",
+    "A delicious ceviche cheesecake slice",
+]
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column():
+        gr.HTML(read_file("html/header.html"))
+    with gr.Column(elem_id="col-container"):
+        input_image = gr.Image(height=800,sources=['upload','clipboard'],image_mode='RGB', elem_id="image_upload", type="pil", label="Upload")
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+                value="A high-resolution photographic image with sharp focus, balanced exposure, clean composition, accurate colour rendering, realistic materials and textures, soft and natural lighting, smooth tonal gradients, minimal noise, high dynamic range, detailed shadows and highlights, precise depth of field, lifelike detail, crisp edges, and visually clear separation between foreground, midground, and background.",
+            )
+
+            run_button = gr.Button("Img2Img", scale=0, variant="primary")
+
+        output_image = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024, 
+                )
+
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                )
+
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=30,
+                    step=1,
+                    value=8,
+                )
+
+        gr.Examples(examples=examples, inputs=[prompt])
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=inference,
+        inputs=[
+            input_image,
+            prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[output_image, seed],
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)