Update app.py

app.py

@@ -6,11 +6,17 @@ import torch
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
 from PIL import Image
 import io
+import os
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
+huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+
+
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token = huggingface_token).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -92,7 +98,8 @@ body {
 }
 """
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo:
+    gr.HTML(title)
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# FLUX.1 [dev]
@@ -105,7 +112,8 @@ with gr.Blocks(css=css) as demo:
 </div>
 </a>
         """)
-        
+
+            
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -225,4 +233,5 @@ with gr.Blocks(css=css) as demo:
         """
     )
 
+
 demo.launch()