| import json | |
| from datetime import datetime | |
| from pathlib import Path | |
| from uuid import uuid4 | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from huggingface_hub import CommitScheduler, InferenceClient | |
| IMAGE_DATASET_DIR = Path("image_dataset") / f"train-{uuid4()}" | |
| IMAGE_DATASET_DIR.mkdir(parents=True, exist_ok=True) | |
| IMAGE_JSONL_PATH = IMAGE_DATASET_DIR / "metadata.jsonl" | |
| scheduler = CommitScheduler( | |
| repo_id="example-space-to-dataset-image", | |
| repo_type="dataset", | |
| folder_path=IMAGE_DATASET_DIR, | |
| path_in_repo=IMAGE_DATASET_DIR.name, | |
| ) | |
| client = InferenceClient() | |
| def generate_image(prompt: str) -> Image: | |
| return client.text_to_image(prompt) | |
| def save_image(prompt: str, image_array: np.ndarray) -> None: | |
| image_path = IMAGE_DATASET_DIR / f"{uuid4()}.png" | |
| with scheduler.lock: | |
| Image.fromarray(image_array).save(image_path) | |
| with IMAGE_JSONL_PATH.open("a") as f: | |
| json.dump({"prompt": prompt, "file_name": image_path.name, "datetime": datetime.now().isoformat()}, f) | |
| f.write("\n") | |
| def get_demo(): | |
| with gr.Row(): | |
| prompt_value = gr.Textbox(label="Prompt") | |
| image_value = gr.Image(label="Generated image") | |
| text_to_image_btn = gr.Button("Generate") | |
| text_to_image_btn.click(fn=generate_image, inputs=prompt_value, outputs=image_value).success( | |
| fn=save_image, | |
| inputs=[prompt_value, image_value], | |
| outputs=None, | |
| ) | |