modified app.py

app.py

@@ -1,3 +1,73 @@
 import gradio as gr
+import torch
+from diffusers import DiffusionPipeline
+import time
 
-gr.load("models/black-forest-labs/FLUX.1-dev").launch()
+# Initialize the base model
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+
+MAX_SEED = 2**32-1
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
+def generate_image(prompt, steps, seed, cfg_scale, width, height):
+    pipe.to("cuda")
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
+    with calculateDuration("Generating image"):
+        # Generate image
+        image = pipe(
+            prompt=prompt,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator
+        ).images[0]
+    return image
+
+def run_model(prompt, cfg_scale, steps, randomize_seed, seed, width, height):
+    if randomize_seed:
+        seed = torch.randint(0, MAX_SEED, (1,)).item()
+    
+    image = generate_image(prompt, steps, seed, cfg_scale, width, height)
+    return image, seed  
+
+with gr.Blocks() as app:
+    gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Type a prompt here")
+            generate_button = gr.Button("Generate")
+    with gr.Row():
+        result = gr.Image(label="Generated Image")
+
+    with gr.Row():
+        with gr.Column():
+            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
+            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
+            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
+            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
+            randomize_seed = gr.Checkbox(True, label="Randomize seed")
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+
+    gr.Interface(
+        fn=run_model,
+        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height],
+        outputs=[result, seed],
+        live=True
+    ).launch()

requirements.txt

@@ -1,6 +1,6 @@
-accelerate
-diffusers
-invisible_watermark
 torch
+git+https://github.com/huggingface/diffusers
+spaces
 transformers
-xformers
+peft
+sentencepiece