File size: 1,446 Bytes
3531f81 1cf80a2 3531f81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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,
)
|