feat: Add FLUX.1 [schnell] model and interface with Gradio

The commit adds the FLUX.1 [schnell] model and integrates it with Gradio for easy generation of images based on user prompts. The code includes the necessary imports, configuration settings, and UI components. It also includes advanced settings for seed, randomizing seed, width, height, and number of inference steps. The commit also includes examples of prompts for generating images.

The commit also includes an update to the requirements.txt file to add the necessary dependencies for Gradio and image loading.

This commit builds upon the previous commits that created the empty app.py file and made the initial commit to the repository.

app.py

@@ -0,0 +1,122 @@
+import gradio as gr
+import numpy as np
+import random
+import spaces
+import torch
+from diffusers import DiffusionPipeline
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1").to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+@spaces.GPU()
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    image = pipe(
+            prompt = prompt, 
+            width = width,
+            height = height,
+            num_inference_steps = num_inference_steps, 
+            generator = generator,
+            guidance_scale=0.0
+    ).images[0] 
+    return image, seed
+ 
+examples = [
+    "a tiny astronaut hatching from an egg on the moon",
+    "a cat holding a sign that says hello world",
+    "an anime illustration of a wiener schnitzel",
+]
+
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# FLUX.1 [schnell]
+12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
+[[blog](https://blackforestlabs.ai/2024/07/31/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-schnell)]
+        """)
+        
+        with gr.Row():
+            
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            
+            run_button = gr.Button("Run", scale=0)
+        
+        result = 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=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,
+                )
+        
+        gr.Examples(
+            examples = examples,
+            fn = infer,
+            inputs = [prompt],
+            outputs = [result, seed],
+            cache_examples="lazy"
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
+        outputs = [result, seed]
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+gradio-imageslider
+loadimg