Update app.py

app.py

@@ -1,25 +1,9 @@
-# Copyright 2024-present, David Berenstein, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import io
 import os
+import io
 import random
-import time
 
 import requests
 from PIL import Image
-
 from dataset_viber import AnnotatorInterFace
 
 HF_TOKEN = os.environ["HF_TOKEN"]
@@ -29,6 +13,10 @@ DATASET_NAME = "poloclub%2Fdiffusiondb&config=2m_random_1k&split=train"
 MODEL_URL = (
     "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 )
+MODEL_URLS = [
+    MODEL_URL,
+    "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+]
 
 
 def retrieve_sample(idx):
@@ -48,17 +36,10 @@ def get_rows():
 
 
 def generate_response(prompt):
-    def _get_response(prompt):
-        payload = {
-            "inputs": prompt,
-        }
-        response = requests.post(MODEL_URL, headers=HEADERS, json=payload)
-        if response.status_code != 200:
-            time.sleep(10)
-            return _get_response(prompt)
-        return response
-
-    response = _get_response(prompt)
+    payload = {
+        "inputs": prompt,
+    }
+    response = requests.post(random.choice(MODEL_URLS), headers=HEADERS, json=payload)
     image = Image.open(io.BytesIO(response.content))
     return image
 
@@ -66,8 +47,7 @@ def generate_response(prompt):
 def next_input(_prompt, _completion_a, _completion_b):
     random_idx = random.randint(0, get_rows()) - 1
     img_url, prompt = retrieve_sample(random_idx)
-    generated_image = generate_response(prompt)
-    return (prompt, img_url, generated_image)
+    return (prompt, generate_response(prompt), generate_response(prompt+" "))
 
 
 if __name__ == "__main__":