Update raw.py

raw.py

@@ -14,6 +14,7 @@ from peft import PeftModel, PeftConfig
 # )
 import gradio as gr
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+MAX_SEED = 1000000
 
 quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
 text_encoder_2_8bit = T5EncoderModel.from_pretrained(
@@ -55,15 +56,19 @@ pipe.unload_lora_weights()
 # pipe.push_to_hub("fused-t-r")
 
 @spaces.GPU
-def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale):
+def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
+    generator = torch.Generator().manual_seed(seed)
+
     # Load control image
     control_image = load_image(control_image)
     w, h = control_image.size
-    # Upscale x1
+    w = w - w % 32
+    h = h - h % 32
     control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
     print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
     with torch.inference_mode():
         image = pipe(
+            generator=generator
             prompt=prompt,
             control_image=control_image,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
@@ -71,6 +76,7 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
             guidance_scale=guidance_scale,
             height=control_image.size[1],
             width=control_image.size[0]
+            control_guidance_end=guidance_end
         ).images[0]
         
     return image
@@ -86,12 +92,14 @@ with gr.Blocks(title="FLUX ControlNet Image Generation", fill_height=True) as if
     with gr.Row():
         with gr.Column(scale=1):
             prompt = gr.Textbox(lines=4, placeholder="Enter your prompt here...", label="Prompt")
+            scale = gr.Slider(1, 3, value=1, label="Scale", step=0.25)
             generate_button = gr.Button("Generate Image", variant="primary")
         with gr.Column(scale=1):
-            scale = gr.Slider(1, 3, value=1, label="Scale")
-            steps = gr.Slider(6, 30, value=8, label="Steps")
-            guidance_scale = gr.Slider(1, 20, value=3.5, label="Guidance Scale")
+            seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
+            steps = gr.Slider(2, 16, value=8, label="Steps")
             controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale")
+            guidance_scale = gr.Slider(1, 20, value=3.5, label="Guidance Scale")
+            guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End")
             
     
     with gr.Row():
@@ -99,7 +107,7 @@ with gr.Blocks(title="FLUX ControlNet Image Generation", fill_height=True) as if
     
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale],
+        inputs=[prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end],
         outputs=[generated_image]
     )