Update optimized.py

optimized.py

@@ -17,6 +17,7 @@ import gc
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 device = "cuda"
 torch_dtype = torch.bfloat16
+MAX_SEED = 1000000
 
 def self_attention_slicing(module, slice_size=3):
     """Modified from Diffusers' original for Flux compatibility"""
@@ -107,7 +108,7 @@ pipe.enable_attention_slicing(1)
 
 print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 @spaces.GPU
-def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, guidance_start, guidance_end):
+def generate_image(prompt, scale, steps, seed, control_image, controlnet_conditioning_scale, guidance_scale, guidance_start, guidance_end):
     print(f"Memory Usage: {torch.cuda.memory_summary(device=None, abbreviated=False)}")
 
     # Load control image
@@ -117,6 +118,8 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
     h = h - h % 8
     control_image = control_image.resize((int(w * scale), int(h * scale)))
     print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
+    generator = torch.Generator().manual_seed(seed)
+
     image = pipe(
         prompt=prompt,
         control_image=control_image,
@@ -126,7 +129,8 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
         height=h,
         width=w,
         control_guidance_start=guidance_start,
-        control_guidance_end=guidance_end
+        control_guidance_end=guidance_end,
+        generator=generator
     ).images[0]
     return image
 # Create Gradio interface
@@ -136,6 +140,7 @@ iface = gr.Interface(
         gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
         gr.Slider(1, 3, value=1, label="Scale"),
         gr.Slider(2, 20, value=8, label="Steps"),
+        gr.Slider(0, MAX_SEED, value=42, label="Seed"),
         gr.Image(type="pil", label="Control Image"),
         gr.Slider(0, 1, value=0.6, label="ControlNet Scale"),
         gr.Slider(1, 20, value=3.5, label="Guidance Scale"),