add prompt confidence

app.py

@@ -154,7 +154,7 @@ def get_random_seed():
 
 # Apply decorator conditionally
 @spaces.GPU() if ENABLE_ZEROGPU else lambda func: func
-def generate_3d(scribble_image_dict, prompt, scribble_confidence, seed): # Added seed parameter back
+def generate_3d(scribble_image_dict, prompt, scribble_confidence, prompt_confidence, seed): # Added text_confidence parameter
     print("Generating 3D model...")
     # Extract the composite image from the ImageEditor dictionary
     if scribble_image_dict is None or scribble_image_dict.get("composite") is None:
@@ -177,7 +177,7 @@ def generate_3d(scribble_image_dict, prompt, scribble_confidence, seed): # Added
     print("Image preprocessed.")
 
     # Define fixed parameters
-    attn_scale_text = 1.0 # As per the example run.py
+    # attn_scale_text = 1.0 # Replaced by text_confidence input
 
     # Set the generator with the provided seed
     generator = torch.Generator(device='cuda').manual_seed(current_seed)
@@ -187,11 +187,11 @@ def generate_3d(scribble_image_dict, prompt, scribble_confidence, seed): # Added
     out = pipe(
         processed_image,
         prompt=prompt,
-        num_tokens=512,
-        guidance_scale=0,
-        num_inference_steps=16,
+        num_tokens=512, # Default value from example
+        guidance_scale=0, # Default value from example
+        num_inference_steps=16, # Default value from example
         attention_kwargs={
-            "cross_attention_scale": attn_scale_text,
+            "cross_attention_scale": prompt_confidence, # Use input parameter
             "cross_attention_2_scale": scribble_confidence
         },
         generator=generator,
@@ -229,6 +229,7 @@ with gr.Blocks() as demo:
             )
             prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., a cute cat wearing a hat")
             confidence_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Scribble Confidence")
+            prompt_confidence_input = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="Prompt Confidence") # Added slider
             seed_input = gr.Number(label="Seed", value=0, precision=0) # Added Seed input back
             with gr.Row():
                 submit_button = gr.Button("Generate 3D Model", variant="primary", scale=1)
@@ -237,16 +238,16 @@ with gr.Blocks() as demo:
             model_output = gr.Model3D(label="Generated 3D Model", interactive=False)
 
     # Define the inputs for the main generation function
-    gen_inputs = [image_input, prompt_input, confidence_input, seed_input]
+    gen_inputs = [image_input, prompt_input, confidence_input, prompt_confidence_input, seed_input] # Added text_confidence_input
 
     submit_button.click(
         fn=generate_3d,
-        inputs=gen_inputs, # Include seed_input
+        inputs=gen_inputs, # Include seed_input and text_confidence_input
         outputs=model_output
     )
 
     # Define inputs for the lucky button (same as main button for the final call)
-    lucky_gen_inputs = [image_input, prompt_input, confidence_input, seed_input]
+    lucky_gen_inputs = [image_input, prompt_input, confidence_input, prompt_confidence_input, seed_input] # Added text_confidence_input
 
     lucky_button.click(
         fn=get_random_seed, # First, get a random seed