Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import torch
+import random
 
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline
 from diffusers import EulerDiscreteScheduler
 
 device = "cpu"
@@ -17,17 +18,24 @@ if mps_available:
     dtype = torch.float16
 #print(f"device: {device}, dtype: {dtype}")
 
-
-pipeline = DiffusionPipeline.from_pretrained("recoilme/ColorfulXL-Lightning",
-                                             variant="fp16",
-                                             torch_dtype=dtype,
-                                             use_safetensors=True)
+pipeline = StableDiffusionXLPipeline.from_pretrained("recoilme/ColorfulXL-Lightning",
+                                            variant="fp16",
+                                            torch_dtype=dtype,
+                                            use_safetensors=True)
 pipeline.to(device)
 pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing")
+# Comes from
+# https://wandb.ai/nasirk24/UNET-FreeU-SDXL/reports/FreeU-SDXL-Optimal-Parameters--Vmlldzo1NDg4NTUw
+if device == "cuda":
+    pipeline.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
+
 
+def generate(prompt, width, height, sample_steps, seed):
+    generator = torch.Generator(device=device).manual_seed(int(seed)) 
+    return pipeline(prompt=prompt, prompt_2=prompt, guidance_scale=0, generator=generator, negative_prompt=None, negative_prompt_2=None, width=width, height=height, num_inference_steps=sample_steps).images[0]
 
-def generate(prompt, width, height, sample_steps):
-    return pipeline(prompt=prompt, guidance_scale=0, negative_prompt="", width=width, height=height, num_inference_steps=sample_steps).images[0]
+def random_seed():
+    return random.randint(0, 2**32 - 1)
 
 with gr.Blocks() as interface:
         with gr.Column():
@@ -41,14 +49,20 @@ with gr.Blocks() as interface:
                             with gr.Column():
                                 width = gr.Slider(label="Width", info="The width in pixels of the generated image.", value=576, minimum=512, maximum=1280, step=64, interactive=True)
                                 height = gr.Slider(label="Height", info="The height in pixels of the generated image.", value=832, minimum=512, maximum=1280, step=64, interactive=True)
+                                with gr.Row():
+                                    seed = gr.Number(label="Seed", 
+                                             value=None,
+                                             scale=8, 
+                                             info="Random seed for reproducibility.")
+                                    seed_button = gr.Button("🎲", scale=2, elem_id="seed_button")
+                                    seed_button.click(fn=random_seed, inputs=[], outputs=seed)
                             with gr.Column():
                                 sampling_steps = gr.Slider(label="Sampling Steps", info="The number of denoising steps.", value=5, minimum=3, maximum=10, step=1, interactive=True)
                         
             with gr.Row():
                 output = gr.Image()
         
-        generate_button.click(fn=generate, inputs=[prompt, width, height, sampling_steps], outputs=[output])
+        generate_button.click(fn=generate, inputs=[prompt, width, height, sampling_steps, seed], outputs=[output])
 
 if __name__ == "__main__":
-    interface.launch()
-
+    interface.launch()