Update app.py

app.py

@@ -1,6 +1,8 @@
 import torch
 import spaces
 import gradio as gr
+import random
+import numpy as np
 from diffusers import ZImagePipeline
 
 # Load the pipeline once at startup
@@ -28,23 +30,21 @@ spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant
 
 pipe_no_lora.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
 spaces.aoti_blocks_load(pipe_no_lora.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
-
+MAX_SEED = np.iinfo(np.int32).max
 print("Pipeline loaded!")
 
 @spaces.GPU
-def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt, height, width, num_inference_steps, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
     """Generate an image from the given prompt."""
-    if randomize_seed:
-        seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    
-    generator = torch.Generator("cuda").manual_seed(int(seed))
+     if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
     image = pipe(
         prompt=prompt,
         height=int(height),
         width=int(width),
         num_inference_steps=int(num_inference_steps),
         guidance_scale=0.0,  # Guidance should be 0 for Turbo models
-        generator=generator,
+        generator = torch.Generator(device=device).manual_seed(seed)
     ).images[0]
 
     image_no_lora = pipe_no_lora(
@@ -53,7 +53,7 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         width=int(width),
         num_inference_steps=int(num_inference_steps),
         guidance_scale=0.0,  # Guidance should be 0 for Turbo models
-        generator=generator,
+        generator = torch.Generator(device=device).manual_seed(seed)
     ).images[0]
     
     return (image_no_lora,image), seed
@@ -124,7 +124,7 @@ with gr.Blocks() as demo:
                     )
                     randomize_seed = gr.Checkbox(
                         label="Randomize Seed",
-                        value=False,
+                        value=True,
                     )
                 
                 generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")