Update app.py

app.py

@@ -1,46 +1,39 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
 import torch
+import spaces
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 
+dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+MAX_IMAGE_SIZE = 2048
 
+@spaces.GPU(duration=190)
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, 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, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
+        width = width,
         height = height,
-        generator = generator
+        num_inference_steps = num_inference_steps, 
+        generator = generator,
+        guidance_scale=guidance_scale
     ).images[0] 
-    
-    return image
-
+    return image, seed
+ 
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    "a cat holding a sign that says hello world",
+    "A scene full of classic video game characters as stickers on a black water bottle",
+    "A futuristic biocity that is located in the former site of Portsmouth, New Hampshire. It has a mix of old and new buildings, green spaces, and water features. It also has six large artificial floating islands off of its coastline,(zenithal angle), ((by Iwan Baan)), coastal city,blue sky and white clouds,the sun is shining brightly,ultra-wide angle,",
+    "Depict a breathtaking scene of a meteor rain showering down from a starry night sky. The meteors should vary in size and brightness, streaking across the sky with vibrant tails of light, creating a dazzling display. Below, a serene landscape—perhaps a tranquil lake reflecting the celestial spectacle, or a rugged mountain range—should enhance the sense of wonder. The foreground can include silhouettes of trees or figures gazing up in awe at the cosmic event. The overall atmosphere should evoke feelings of magic and inspiration, capturing the beauty and mystery of the universe.",
 ]
 
 css="""
@@ -50,17 +43,12 @@ css="""
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
+        gr.Markdown(f"""# FLUX.1 [dev]
+12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
+[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
         """)
         
         with gr.Row():
@@ -76,16 +64,9 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0)
         
         result = gr.Image(label="Result", show_label=False)
-
+        
         with gr.Accordion("Advanced Settings", open=False):
             
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -103,7 +84,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
                 
                 height = gr.Slider(
@@ -111,36 +92,40 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
             
             with gr.Row():
-                
+
                 guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=15,
                     step=0.1,
-                    value=0.0,
+                    value=3.5,
                 )
-                
+  
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
-                    maximum=12,
+                    maximum=50,
                     step=1,
-                    value=2,
+                    value=28,
                 )
         
         gr.Examples(
             examples = examples,
-            inputs = [prompt]
+            fn = infer,
+            inputs = [prompt],
+            outputs = [result, seed],
+            cache_examples="lazy"
         )
 
-    run_button.click(
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs = [result, seed]
     )
 
-demo.queue().launch()
+demo.launch()