Update app.py

app.py

@@ -1,66 +1,110 @@
 import gradio as gr
+import numpy as np
+import random
 import torch
+import time
 from diffusers import DiffusionPipeline
 
-# Check for CUDA availability
+# Set the device and dtype
+dtype = torch.float16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-if device == "cpu":
-    print("Warning: CUDA is not available. This model may not run correctly on CPU.")
 
-# Load the model
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
-pipe = pipe.to(device)
+# Load the diffusion pipeline
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=7.5, progress=gr.Progress(track_tqdm=True)):
+    start_time = time.time()
+
+    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
+        raise ValueError("Image size exceeds the maximum allowed dimensions.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-def infer(prompt, negative_prompt="", width=1024, height=1024, num_inference_steps=4, guidance_scale=7.5):
     try:
-        # Generate the image
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
             height=height,
             num_inference_steps=num_inference_steps,
+            generator=generator,
             guidance_scale=guidance_scale
         ).images[0]
-        
-        return image
     except Exception as e:
-        return str(e)
+        print(f"Error generating image: {e}")
+        return None, seed, f"Error: {str(e)}"
+
+    if time.time() - start_time > 60:
+        return None, seed, "Image generation took too long and was cancelled."
+
+    return image, seed, None
+
+examples = [
+    ["a tiny astronaut hatching from an egg on the moon", "blurry, low quality"],
+    ["a cat holding a sign that says hello world", "dog, text, writing"],
+    ["an anime illustration of a wiener schnitzel", "realistic, photograph"],
+]
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # Custom Image Creator
+    12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
+    [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1)]
+    """)
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Image Generation with FLUX.1-schnell")
-    
     with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt")
-            negative_prompt = gr.Textbox(label="Negative Prompt")
-            generate_btn = gr.Button("Generate")
-        
-        image_output = gr.Image(label="Generated Image")
-    
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your prompt",
+                lines=3
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="Enter things to avoid in the image",
+                lines=2
+            )
+            run_button = gr.Button("Generate Image", variant="primary")
+
+        with gr.Column(scale=2):
+            result = gr.Image(label="Generated Image")
+            seed_output = gr.Number(label="Seed Used")
+
     with gr.Accordion("Advanced Settings", open=False):
-        width = gr.Slider(minimum=256, maximum=1024, value=1024, step=32, label="Width")
-        height = gr.Slider(minimum=256, maximum=1024, value=1024, step=32, label="Height")
-        num_inference_steps = gr.Slider(minimum=1, maximum=50, value=4, step=1, label="Number of Inference Steps")
-        guidance_scale = gr.Slider(minimum=0, maximum=20, value=7.5, step=0.1, label="Guidance Scale")
-
-    generate_btn.click(
-        infer,
-        inputs=[prompt, negative_prompt, width, height, num_inference_steps, guidance_scale],
-        outputs=[image_output]
-    )
+        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
+        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        with gr.Row():
+            width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+            height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+
+        with gr.Row():
+            num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4)
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.5, value=7.5)
 
     gr.Examples(
-        examples=[
-            ["a cat sitting on a moon", "dog, low quality"],
-            ["a futuristic city skyline", "old buildings, cars"],
-        ],
+        examples=examples,
         inputs=[prompt, negative_prompt],
-        outputs=[image_output],
+        outputs=[result, seed_output],
         fn=infer,
         cache_examples=True
     )
 
+    run_button.click(
+        fn=infer,
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
+        outputs=[result, seed_output]
+    )
+
+    gr.Markdown("""
+    ## Save Your Image
+    Right-click on the generated image and select 'Save image as' to download it.
+    """)
+
 if __name__ == "__main__":
     demo.launch()