Update app.py

app.py

@@ -89,12 +89,12 @@ def generate_image(
         ).images[0]
     
     print(neg_json_prompt)
-    return image, seed, json_prompt, json.dumps(neg_json_prompt), gr.update(visible=True)
+    return image, seed, json_prompt, json.dumps(neg_json_prompt), gr.update(visible=True), image, json_prompt, seed, seed
 
 @spaces.GPU(duration=300)
 def refine_prompt(
     refine_instruction,
-    prompt_in_json,
+    refine_json,
     negative_prompt="",
     seed=42,
     randomize_seed=False,
@@ -107,10 +107,17 @@ def refine_prompt(
         seed = random.randint(0, MAX_SEED)
 
     with torch.inference_mode():
+        # Parse the JSON string if it's a string
+        if isinstance(refine_json, str):
+            try:
+                refine_json = json.loads(refine_json) if refine_json else {}
+            except:
+                refine_json = {}
+        
         json_prompt_str = (
-            json.dumps(prompt_in_json)
-            if isinstance(prompt_in_json, (dict, list))
-            else str(prompt_in_json)
+            json.dumps(refine_json)
+            if isinstance(refine_json, (dict, list))
+            else str(refine_json)
         )
         output = vlm_pipe(json_prompt=json_prompt_str, prompt=refine_instruction)
         json_prompt = output.values["json_prompt"]
@@ -131,7 +138,7 @@ def refine_prompt(
         ).images[0]
     
     print(neg_json_prompt)
-    return image, seed, json_prompt, json.dumps(neg_json_prompt)
+    return image, seed, json_prompt, json.dumps(neg_json_prompt), image, json_prompt, seed
 
 
 css = """
@@ -155,33 +162,60 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
 [[non-commercial license]()] [[arxiv](https://arxiv.org/)] [[model](https://huggingface.co/briaai/FIBO)] [[code](https://github.com/Bria-AI/FIBO)]
         """ , elem_id="app-title")
 
+    # State to store the last used seed
+    last_seed_state = gr.State(value=0)
+
     with gr.Row(elem_id="col-container"):
         with gr.Column(scale=1):
             
-            with gr.Accordion("Inspire from Image Reference", open=False) as inspire_accordion:
-                prompt_inspire_image = gr.Image(
-                    label="Inspiration Image",
-                    type="pil",
-                )
-            
-            prompt_generate = gr.Textbox(
-                label="Prompt",
-                placeholder="a man holding a goose screaming",
-                lines=3
-            )
-            
-            prompt_in_json = gr.JSON(
-                label="Structured JSON Prompt (editable)",
-                value=None
-            )
-            
-            with gr.Row():
-                generate_json_btn = gr.Button("Generate JSON Prompt", variant="secondary")
-                generate_image_btn = gr.Button("Generate New Image", variant="primary")
+            with gr.Tabs() as tabs:
+                with gr.Tab("Generate", id="generate_tab") as tab_generate:
+                    with gr.Accordion("Inspire from Image", open=False) as inspire_accordion:
+                        prompt_inspire_image = gr.Image(
+                            label="Inspiration Image",
+                            type="pil",
+                        )
+                    
+                    prompt_generate = gr.Textbox(
+                        label="Prompt",
+                        placeholder="a man holding a goose screaming",
+                        lines=3
+                    )
+                    
+                    prompt_in_json = gr.Code(
+                        label="Structured JSON Prompt (editable)",
+                        language="json",
+                        value="",
+                        lines=10
+                    )
+                    
+                    with gr.Row():
+                        generate_json_btn = gr.Button("Generate JSON Prompt", variant="secondary")
+                        generate_image_btn = gr.Button("Generate Image", variant="primary")
+                
+                with gr.Tab("Refine", id="refine_tab", visible=False) as tab_refine:
+                    previous_result = gr.Image(label="Previous Generation", interactive=False)
+                    
+                    refine_instruction = gr.Textbox(
+                        label="Refinement instruction",
+                        placeholder="make the cat white",
+                        lines=2,
+                        info="Describe the changes you want to make"
+                    )
+                    
+                    refine_json = gr.Code(
+                        label="Structured JSON Prompt (editable)",
+                        language="json",
+                        value="",
+                        lines=10
+                    )
+                    
+                    with gr.Row():
+                        refine_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                        refine_randomize_seed = gr.Checkbox(label="Randomize seed", value=False, info="Keep off for consistent refinements")
+                    
+                    refine_btn = gr.Button("Refine Image", variant="primary")
             
-            # Refine button - initially hidden
-            refine_btn = gr.Button("Refine Existing Image", variant="primary", visible=False)
-
             with gr.Accordion("Advanced Settings", open=False):
                 with gr.Row():
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
@@ -195,7 +229,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
                     width = gr.Slider(label="Width", minimum=832, maximum=1344, step=64, value=1024)
                 with gr.Row():
                     negative_prompt = gr.Textbox(label="negative prompt")
-                    negative_prompt_json = gr.JSON(label="json negative prompt")
+                    negative_prompt_json = gr.Code(label="json negative prompt", language="json")
 
         with gr.Column(scale=1):
             result = gr.Image(label="output")
@@ -225,24 +259,28 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed, prompt_in_json, negative_prompt_json, refine_btn],
+        outputs=[result, seed, prompt_in_json, negative_prompt_json, tab_refine, previous_result, refine_json, last_seed_state],
+    ).then(
+        fn=lambda s: s,  # Copy seed to refine tab
+        inputs=[last_seed_state],
+        outputs=[refine_seed]
     )
 
-    # Refine image (reuses the main prompt box)
+    # Refine image with dedicated seed controls
     refine_btn.click(
         fn=refine_prompt,
         inputs=[
-            prompt_generate,  # Reuse the main prompt box
-            prompt_in_json,
+            refine_instruction,
+            refine_json,
             negative_prompt,
-            seed,
-            randomize_seed,
+            refine_seed,
+            refine_randomize_seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed, prompt_in_json, negative_prompt_json],
+        outputs=[result, refine_seed, refine_json, negative_prompt_json, previous_result, refine_json, last_seed_state],
     )
 
 demo.queue().launch()