Update app.py

app.py

@@ -9,6 +9,7 @@ import numpy as np
 from diffusers import BriaFiboPipeline
 from diffusers.modular_pipelines import ModularPipeline
 import requests
+import io
 
 MAX_SEED = np.iinfo(np.int32).max
 dtype = torch.bfloat16
@@ -44,25 +45,33 @@ def get_default_negative_prompt(existing_json: dict) -> str:
 
 def generate_json_prompt(
     prompt,
+    json_prompt,
     prompt_inspire_image=None,
     seed=42
 ):
-    url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate/pro"
-    
+
     api_key = os.environ.get("BRIA_API_TOKEN")
     
-    payload = {
-    "prompt": prompt,
-    "seed": seed,
-    "sync": True
-  }
-    headers = {
-    "Content-Type": "application/json",
-    "api_token": api_key
-  }
+    url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate/pro"
+
+    payload = {"seed": seed, "sync": True}
+    if json_prompt:
+        payload["structured_prompt"] = structured_prompt
+    if prompt:
+        payload["prompt"] = prompt
+    if prompt_inspire_image:
+        buffered = io.BytesIO()
+        prompt_inspire_image.save(buffered, format="PNG")
+        image_bytes = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        payload["images"] = [image_bytes]
+
+    headers = {"Content-Type": "application/json", "api_token": api_key}
+
     response = requests.post(url, json=payload, headers=headers)
+
     data = response.json()
-    return data['result']['structured_prompt']
+
+    return data["result"]["structured_prompt"]
 
 @spaces.GPU(duration=300)
 def generate_image(
@@ -86,7 +95,7 @@ def generate_image(
         if not prompt_in_json or prompt_in_json == "":
             if prompt_inspire_image is not None:
                 #output = vlm_pipe(image=prompt_inspire_image, prompt="")
-                json_prompt = generate_json_prompt(prompt=prompt, seed=seed)
+                json_prompt = generate_json_prompt(prompt=prompt, prompt_inspire_image=prompt_inspire_image, seed=seed)
             else:
                 # output = vlm_pipe(prompt=prompt)
                 json_prompt = generate_json_prompt(prompt=prompt, seed=seed)
@@ -147,12 +156,14 @@ def refine_prompt(
             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"]
+        # output = vlm_pipe(json_prompt=json_prompt_str, prompt=refine_instruction)
+        # json_prompt = output.values["json_prompt"]
+        json_prompt = generate_json_prompt(prompt=refine_instruction, json_prompt=json_prompt_str, seed=seed)
 
         if negative_prompt:
-            neg_output = vlm_pipe(prompt=negative_prompt)
-            neg_json_prompt = neg_output.values["json_prompt"]
+            # neg_output = vlm_pipe(prompt=negative_prompt)
+            # neg_json_prompt = neg_output.values["json_prompt"]
+            neg_json_prompt = generate_json_prompt(prompt=negative_prompt, seed=seed)
         else:
             neg_json_prompt = get_default_negative_prompt(json.loads(json_prompt))