| import os | |
| import io | |
| import random | |
| import requests | |
| from PIL import Image | |
| from data_viber import AnnotatorInterFace | |
| HF_TOKEN = os.environ["HF_TOKEN"] | |
| HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} | |
| DATASET_SERVER_URL = "https://datasets-server.huggingface.co" | |
| DATASET_NAME = "poloclub%2Fdiffusiondb&config=2m_random_1k&split=train" | |
| MODEL_URL = ( | |
| "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell" | |
| ) | |
| def retrieve_sample(idx): | |
| api_url = f"{DATASET_SERVER_URL}/rows?dataset={DATASET_NAME}&offset={idx}&length=1" | |
| response = requests.get(api_url, headers=HEADERS) | |
| data = response.json() | |
| img_url = data["rows"][0]["row"]["image"]["src"] | |
| prompt = data["rows"][0]["row"]["prompt"] | |
| return img_url, prompt | |
| def get_rows(): | |
| api_url = f"{DATASET_SERVER_URL}/size?dataset={DATASET_NAME}" | |
| response = requests.get(api_url, headers=HEADERS) | |
| num_rows = response.json()["size"]["config"]["num_rows"] | |
| return num_rows | |
| def generate_response(prompt): | |
| payload = { | |
| "inputs": prompt, | |
| } | |
| response = requests.post(MODEL_URL, headers=HEADERS, json=payload) | |
| image = Image.open(io.BytesIO(response.content)) | |
| return image | |
| 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) | |
| if __name__ == "__main__": | |
| interface = AnnotatorInterFace.for_image_generation_preference( | |
| fn=next_input, | |
| dataset_name=None, | |
| ) | |
| interface.launch() | |