Update app.py

app.py

@@ -14,9 +14,18 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-vlm_pipe = ModularPipeline.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True).to(device)
+vlm_pipe = ModularPipeline.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True)
+# vlm_pipe = ModularPipeline.from_pretrained("briaai/FIBO-gemini-prompt-to-JSON", trust_remote_code=True)
 pipe = BriaFiboPipeline.from_pretrained("briaai/FIBO", trust_remote_code=True, torch_dtype=dtype).to(device)
 
+@spaces.GPU()
+def get_default_negative_prompt(existing_json: dict) -> str:
+    negative_prompt = ""
+    style_medium = existing_json.get("style_medium", "").lower()
+    if style_medium in ["photograph", "photography", "photo"]:
+        negative_prompt = """{'style_medium':'digital illustration','artistic_style':'non-realistic'}"""
+    return negative_prompt
+
 @spaces.GPU(duration=300)
 def infer(
     prompt,
@@ -26,7 +35,7 @@ def infer(
     seed=42,
     randomize_seed=False,
     width=1024,
-    height=768,
+    height=1024,
     guidance_scale=5,
     num_inference_steps=50,
     mode="generate",
@@ -35,12 +44,7 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     with torch.inference_mode():
-        if negative_prompt:
-            neg_output = vlm_pipe(prompt=negative_prompt)
-            neg_json_prompt = neg_output.values["json_prompt"]
-        else:
-            neg_json_prompt = ""
-            
+         
         if mode == "refine":
             json_prompt_str = (
                 json.dumps(prompt_in_json)
@@ -52,6 +56,12 @@ def infer(
             output = vlm_pipe(prompt=prompt)
         json_prompt = output.values["json_prompt"]
 
+        if negative_prompt:
+            neg_output = vlm_pipe(prompt=negative_prompt)
+            neg_json_prompt = neg_output.values["json_prompt"]
+        else:
+            neg_json_prompt = get_default_negative_prompt(json.loads(json_prompt))
+
         image = pipe(
             prompt=json_prompt,
             num_inference_steps=num_inference_steps,