Update app.py

app.py

@@ -30,6 +30,7 @@ def get_default_negative_prompt(existing_json: dict) -> str:
 def infer(
     prompt,
     prompt_refine,
+    prompt_inspire_image,
     prompt_in_json,
     negative_prompt="",
     seed=42,
@@ -52,6 +53,12 @@ def infer(
                 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="")
+
         else:
             output = vlm_pipe(prompt=prompt)
         json_prompt = output.values["json_prompt"]
@@ -70,7 +77,7 @@ def infer(
             height=height,
             guidance_scale=guidance_scale,
         ).images[0]
-
+    print(neg_json_prompt)
     return image, seed, json_prompt, neg_json_prompt
 
 
@@ -94,6 +101,14 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
                 with gr.Row():
                     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:
+                with gr.Row():
+                    prompt_inspire_image = gr.Image(
+                        label="Inspiration Image",
+                        type="pil",
+                        tool=None
+                    )
+            
         submit_btn = gr.Button("Generate", variant="primary")
         result = gr.Image(label="output")
         with gr.Accordion("Structured Prompt", open=False):
@@ -118,12 +133,14 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
 
         tab_generate.select(lambda: "generate", outputs=current_mode)
         tab_refine.select(lambda: "refine", outputs=current_mode)
+        tab_inspire.select(lambda: "inspire", outputs=current_mode)
 
         submit_btn.click(
             fn=infer,
             inputs=[
                 prompt_generate,
                 prompt_refine,
+                prompt_inspire_image,
                 prompt_in_json,
                 negative_prompt,
                 seed,