Update app.py

app.py

@@ -27,10 +27,26 @@ def get_default_negative_prompt(existing_json: dict) -> str:
     return negative_prompt
 
 @spaces.GPU(duration=300)
-def infer(
+def generate_json_prompt(
     prompt,
-    prompt_refine,
     prompt_inspire_image,
+    use_inspire,
+):
+    with torch.inference_mode():
+        if use_inspire and prompt_inspire_image is not None:
+            output = vlm_pipe(image=prompt_inspire_image, prompt="")
+        else:
+            output = vlm_pipe(prompt=prompt)
+        
+        json_prompt = output.values["json_prompt"]
+    
+    return json_prompt
+
+@spaces.GPU(duration=300)
+def generate_image(
+    prompt,
+    prompt_inspire_image,
+    use_inspire,
     prompt_in_json,
     negative_prompt="",
     seed=42,
@@ -39,28 +55,66 @@ def infer(
     height=1024,
     guidance_scale=5,
     num_inference_steps=50,
-    mode="generate",
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     with torch.inference_mode():
-         
-        if mode == "refine":
-            json_prompt_str = (
+        # If JSON prompt is empty or None, generate it first
+        if not prompt_in_json or prompt_in_json == "":
+            if use_inspire and prompt_inspire_image is not None:
+                output = vlm_pipe(image=prompt_inspire_image, prompt="")
+            else:
+                output = vlm_pipe(prompt=prompt)
+            json_prompt = output.values["json_prompt"]
+        else:
+            # Use the provided JSON prompt
+            json_prompt = (
                 json.dumps(prompt_in_json)
                 if isinstance(prompt_in_json, (dict, list))
                 else str(prompt_in_json)
             )
-            output = vlm_pipe(json_prompt=json_prompt_str, prompt=prompt_refine)
-            
-        elif mode == "inspire":
-            if prompt_inspire_image is None:
-                raise ValueError("Please upload an image to inspire the model.")
-            output = vlm_pipe(image=prompt_inspire_image, prompt="")
 
+        if negative_prompt:
+            neg_output = vlm_pipe(prompt=negative_prompt)
+            neg_json_prompt = neg_output.values["json_prompt"]
         else:
-            output = vlm_pipe(prompt=prompt)
+            neg_json_prompt = get_default_negative_prompt(json.loads(json_prompt))
+
+        image = pipe(
+            prompt=json_prompt,
+            num_inference_steps=num_inference_steps,
+            negative_prompt=neg_json_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+        ).images[0]
+    
+    print(neg_json_prompt)
+    return image, seed, json_prompt, json.dumps(neg_json_prompt), gr.update(visible=True)
+
+@spaces.GPU(duration=300)
+def refine_prompt(
+    refine_instruction,
+    prompt_in_json,
+    negative_prompt="",
+    seed=42,
+    randomize_seed=False,
+    width=1024,
+    height=1024,
+    guidance_scale=5,
+    num_inference_steps=50,
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    with torch.inference_mode():
+        json_prompt_str = (
+            json.dumps(prompt_in_json)
+            if isinstance(prompt_in_json, (dict, list))
+            else str(prompt_in_json)
+        )
+        output = vlm_pipe(json_prompt=json_prompt_str, prompt=refine_instruction)
         json_prompt = output.values["json_prompt"]
 
         if negative_prompt:
@@ -77,8 +131,9 @@ def infer(
             height=height,
             guidance_scale=guidance_scale,
         ).images[0]
+    
     print(neg_json_prompt)
-    return image, seed, json_prompt, json.dumps(neg_json_prompt), gr.update(open=True)
+    return image, seed, json_prompt, json.dumps(neg_json_prompt)
 
 
 css = """
@@ -104,28 +159,31 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
 
     with gr.Row(elem_id="col-container"):
         with gr.Column(scale=1):
-
+            
+            with gr.Accordion("Inspire from Image", open=False) as inspire_accordion:
+                prompt_inspire_image = gr.Image(
+                    label="Inspiration Image",
+                    type="pil",
+                )
+                use_inspire = gr.Checkbox(label="Use inspiration image", value=False)
+            
+            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():
-                with gr.Tab("generate") as tab_generate:
-                    prompt_generate = gr.Textbox(
-                        label="Prompt",
-                        placeholder="a man holding a goose screaming"
-                    )
-
-                with gr.Tab("refine") as tab_refine:
-                    prompt_refine = gr.Textbox(
-                        label="Prompt",
-                        info="describe the change you want to make",
-                        placeholder="make the cat white"
-                    )
-
-                with gr.Tab("inspire") as tab_inspire:
-                    prompt_inspire_image = gr.Image(
-                        label="Inspiration Image",
-                        type="pil",
-                    )
-
-            submit_btn = gr.Button("Generate", variant="primary")
+                generate_json_btn = gr.Button("Generate JSON Prompt", variant="secondary")
+                generate_image_btn = gr.Button("Generate 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():
@@ -144,27 +202,42 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
 
         with gr.Column(scale=1):
             result = gr.Image(label="output")
-            with gr.Accordion("Structured Prompt", open=False) as structured_accordion:
-                prompt_in_json = gr.JSON(label="json structured prompt")
-
-    # Track active tab
-    current_mode = gr.State("generate")
-    # When "generate" is selected — just set mode
-    tab_generate.select(lambda: ("generate", gr.update(value=True)), outputs=[current_mode, randomize_seed])
-    
-    # When "refine" is selected — set mode and turn off randomize_seed
-    tab_refine.select(lambda: ("refine", gr.update(value=False)), outputs=[current_mode, randomize_seed])
-    
-    # When "inspire" is selected — normal
-    tab_inspire.select(lambda: ("inspire", gr.update(value=True)), outputs=[current_mode, randomize_seed])
 
+    # Generate JSON prompt only
+    generate_json_btn.click(
+        fn=generate_json_prompt,
+        inputs=[
+            prompt_generate,
+            prompt_inspire_image,
+            use_inspire,
+        ],
+        outputs=[prompt_in_json],
+    )
 
-    submit_btn.click(
-        fn=infer,
+    # Generate image (generates JSON first if needed)
+    generate_image_btn.click(
+        fn=generate_image,
         inputs=[
             prompt_generate,
-            prompt_refine,
             prompt_inspire_image,
+            use_inspire,
+            prompt_in_json,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[result, seed, prompt_in_json, negative_prompt_json, refine_btn],
+    )
+
+    # Refine image (reuses the main prompt box)
+    refine_btn.click(
+        fn=refine_prompt,
+        inputs=[
+            prompt_generate,  # Reuse the main prompt box
             prompt_in_json,
             negative_prompt,
             seed,
@@ -173,9 +246,8 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
             height,
             guidance_scale,
             num_inference_steps,
-            current_mode,
         ],
-        outputs=[result, seed, prompt_in_json, negative_prompt_json, structured_accordion],
+        outputs=[result, seed, prompt_in_json, negative_prompt_json],
     )
 
-demo.queue().launch()
+demo.queue().launch()