Update app.py

app.py

@@ -298,8 +298,43 @@ def numpy_to_list(array):
 #     pipe=load_model()
 #     image = pipe(prompt).images[0]
 #     return image
+def refer_api(prompt):
+    invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/stable-diffusion-3-medium"
 
+    headers = {
+        "Authorization": "Bearer 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",
+        "Accept": "application/json",
+    }
+    
+    payload = {
+        "prompt": prompt,
+        "cfg_scale": 5,
+        "aspect_ratio": "16:9",
+        "seed": 0,
+        "steps": 50,
+        "negative_prompt": ""
+    }
+    
+
+    response = requests.post(invoke_url, headers=headers, json=payload)
+    
+    response.raise_for_status()
+    response_body = response.json()
+    def string_to_image(base64_string):
+        if ',' in base64_string:
+            base64_string = base64_string.split(',')[1]
+        # Decode base64 string to bytes
+        image_data = base64.b64decode(base64_string)
+        # Open image from bytes
+        image = Image.open(BytesIO(image_data))
+        return image
+    return string_to_image(response_body['image'])
+    
 def model_out_put(init_image,mask_image,prompt,negative_prompt):
+    if all(mask_image):
+        l,m=init_image.shape
+        image =refer_api(prompt).resize((l,m))
+        return image
     API_URL = "https://8417-201-238-124-65.ngrok-free.app/api/llm-response"
     initial_image_base64 = numpy_to_list(np.array(init_image))
     mask_image_base64 = numpy_to_list(np.array(mask_image))