Update app.py

app.py

@@ -8,10 +8,7 @@ import spaces
 from src.pipeline import DiT360Pipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 model_repo = "black-forest-labs/FLUX.1-dev"
 lora_weights = "fenghora/DiT360-Panorama-Image-Generation"
@@ -26,17 +23,13 @@ MAX_IMAGE_SIZE = 2048
 def infer(
     prompt,
     seed,
-    randomize_seed,
     width,
     guidance_scale,
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
     height = width // 2
-    generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(int(seed))
 
     full_prompt = f"This is a panorama. The images shows {prompt.strip()}"
 
@@ -50,8 +43,10 @@ def infer(
     ).images[0]
     image.save("test.png")
 
-    return image, seed
+    return image
 
+def generate_seed():
+    return random.randint(0, MAX_SEED)
 
 examples = [
     "A medieval castle stands proudly on a hilltop surrounded by autumn forests, with golden light spilling across the landscape.",
@@ -60,7 +55,6 @@ examples = [
     "A snowy mountain village under northern lights, with cozy cabins and smoke rising from chimneys.",
 ]
 
-
 css = """
 #main-container {
     display: flex;
@@ -70,24 +64,20 @@ css = """
     gap: 2rem;
     margin-top: 1rem;
 }
-
 #image-panel {
-    flex: 2; /* 占2/3 */
+    flex: 2;
     max-width: 900px;
     margin: 0 auto;
 }
-
 #settings-panel {
-    flex: 1; /* 占1/3 */
+    flex: 1;
     max-width: 280px;
 }
-
 #prompt-box textarea {
-    resize: none !important; /* 去掉上下箭头 */
+    resize: none !important;
 }
 """
 
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# 🌀 DiT360: High-Fidelity Panoramic Image Generation")
     gr.Markdown("Official Gradio demo for **DiT360**, a panoramic image generation model based on hybrid training.")
@@ -115,10 +105,10 @@ with gr.Blocks(css=css) as demo:
                 "The height is automatically set to half the width (2:1 aspect ratio)."
             )
 
-            seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+            seed_display = gr.Number(value=0, label="Seed", interactive=True)
+            random_seed_button = gr.Button("🎲 Random Seed")
 
-            width = gr.Slider(1024, MAX_IMAGE_SIZE, value=2048, step=64, label="Width (fixed 2:1)")
+            width = gr.Slider(1024, MAX_IMAGE_SIZE, value=2048, step=64, label="Width")
             height_display = gr.Number(value=1024, label="Height", interactive=False)
 
             guidance_scale = gr.Slider(0.0, 10.0, value=2.8, step=0.1, label="Guidance Scale")
@@ -128,6 +118,7 @@ with gr.Blocks(css=css) as demo:
                 return width // 2
 
             width.change(fn=update_height, inputs=width, outputs=height_display)
+            random_seed_button.click(fn=generate_seed, inputs=[], outputs=seed_display)
 
     gr.Markdown(
         "💡 *Tip: Try descriptive prompts like “A mountain village at sunrise with mist over the valley.” "
@@ -137,10 +128,9 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, seed, randomize_seed, width, guidance_scale, num_inference_steps],
-        outputs=[result, seed],
+        inputs=[prompt, seed_display, width, guidance_scale, num_inference_steps],
+        outputs=[result],
     )
 
-
 if __name__ == "__main__":
     demo.launch()