Update app_parquet.py

app_parquet.py

@@ -187,16 +187,16 @@ PARQUET_DATASET_DIR.mkdir(parents=True, exist_ok=True)
 scheduler = ParquetScheduler(repo_id="example-space-to-dataset-parquet")
 
 # client = Client("stabilityai/stable-diffusion") # Space is paused
-client = Client("runwayml/stable-diffusion-v1-5")
+# client = Client("runwayml/stable-diffusion-v1-5") # Space has been deleted
+client = Client("black-forest-labs/FLUX.1-schnell")
 
 
 def generate(prompt: str) -> tuple[str, list[str]]:
     """Generate images on 'submit' button."""
     # Generate from https://huggingface.co/spaces/stabilityai/stable-diffusion
     # out_dir = client.predict(prompt, "", 9, fn_index=1) # Space 'stabilityai/stable-diffusion' is paused
-    out_dir = client.predict(prompt, fn_index=1)
-    with (Path(out_dir) / "captions.json").open() as f:
-        paths = list(json.load(f).keys())
+    # out_dir = client.predict(prompt, fn_index=1) # Space "runwayml/stable-diffusion-v1-5" has been deleted
+    image_path, _ = client.predict(prompt, api_name="/infer")
 
     # Save config used to generate data
     with tempfile.NamedTemporaryFile(
@@ -206,7 +206,7 @@ def generate(prompt: str) -> tuple[str, list[str]]:
             {"prompt": prompt, "negative_prompt": "", "guidance_scale": 9}, config_file
         )
 
-    return config_file.name, paths
+    return config_file.name, image_path
 
 
 def get_selected_index(evt: gr.SelectData) -> int:
@@ -215,7 +215,7 @@ def get_selected_index(evt: gr.SelectData) -> int:
 
 
 def save_preference(
-    config_path: str, gallery: list[dict[str, Any]], selected_index: int
+    config_path: str, image_path: str
 ) -> None:
     """Save preference, i.e. move images to a new folder and send paths+config to scheduler."""
     save_dir = PARQUET_DATASET_DIR / f"{uuid.uuid4()}"
@@ -226,14 +226,14 @@ def save_preference(
         data = json.load(f)
 
     # Add selected item + timestamp
-    data["selected_index"] = selected_index
+    data["selected_index"] = 0
     data["timestamp"] = datetime.datetime.utcnow().isoformat()
 
     # Copy and add images
-    for index, path in enumerate(x["name"] for x in gallery):
+    for index in range(4): # fake 4 images
         name = f"{index:03d}"
-        dst_path = save_dir / f"{name}{Path(path).suffix}"
-        shutil.move(path, dst_path)
+        dst_path = save_dir / f"{name}{Path(image_path).suffix}"
+        shutil.copyfile(image_path, dst_path)
         data[f"image_{name}"] = dst_path
 
     # Send to scheduler
@@ -261,14 +261,9 @@ def get_demo():
     )
 
     # Save preference on click
-    gallery.select(
-        fn=get_selected_index,
-        outputs=selected_index,
-        queue=False,
-    )
     save_preference_button.click(
         fn=save_preference,
-        inputs=[config_path, gallery, selected_index],
+        inputs=[config_path, gallery],
         queue=False,
     ).then(
         fn=clear,