Update app.py

app.py

@@ -197,7 +197,7 @@ class CustomFluxAttnProcessor2_0:
 print("--- Loading Models: This may take a few minutes and requires >40GB VRAM ---")
 controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
 pipe = FluxControlNetInpaintingPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", controlnet=controlnet, torch_dtype=torch.bfloat16).to(device)
-pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+# pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
 
 pipe.transformer.to(torch.bfloat16)
 pipe.controlnet.to(torch.bfloat16)
@@ -211,9 +211,27 @@ segment_processor = AutoProcessor.from_pretrained(segmenter_id)
 object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=device)
 print("--- All models loaded successfully! ---")
 
+def get_duration(
+    input_image: Image.Image,
+    subject_name: str,
+    target_prompt: str,
+    attn_enforce: float = 1.3,
+    ctrl_scale: float = 0.95,
+    width: int = 1024,
+    height: int = 1024,
+    pixel_offset: int = 8,
+    num_steps: int = 28,
+    guidance: float = 3.5,
+    seed: int = 42,
+    randomize_seed: bool = False,
+    progress=gr.Progress(track_tqdm=True)
+):
+    if width > 768 and height > 768:
+        return 210
+    else:
+        return 120
 
-# --- Main Inference Function for Gradio ---
-@spaces.GPU(duration=180)
+@spaces.GPU(duration=get_duration)
 def run_diptych_prompting(
     input_image: Image.Image,
     subject_name: str,
@@ -223,7 +241,7 @@ def run_diptych_prompting(
     width: int = 1024,
     height: int = 1024,
     pixel_offset: int = 8,
-    num_steps: int = 8,
+    num_steps: int = 28,
     guidance: float = 3.5,
     seed: int = 42,
     randomize_seed: bool = False,
@@ -308,7 +326,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             with gr.Accordion("Advanced Settings", open=False):
                 attn_enforce = gr.Slider(minimum=1.0, maximum=2.0, value=1.3, step=0.05, label="Attention Enforcement")
                 ctrl_scale = gr.Slider(minimum=0.5, maximum=1.0, value=0.95, step=0.01, label="ControlNet Scale")
-                num_steps = gr.Slider(minimum=20, maximum=50, value=8, step=1, label="Inference Steps")
+                num_steps = gr.Slider(minimum=20, maximum=50, value=28, step=1, label="Inference Steps")
                 guidance = gr.Slider(minimum=1.0, maximum=10.0, value=3.5, step=0.1, label="Guidance Scale")
                 width = gr.Slider(minimum=512, maximum=1024, value=768, step=64, label="Image Width")
                 height = gr.Slider(minimum=512, maximum=1024, value=768, step=64, label="Image Height")