Upload folder using huggingface_hub

README.md

@@ -4,10 +4,12 @@ emoji: 🌍
 colorFrom: indigo
 colorTo: red
 sdk: gradio
-sdk_version: 3.39.0
+sdk_version: 3.44.2
 app_file: app.py
 pinned: false
 license: mit
+duplicated_from: warp-ai/Wuerstchen
+hf_oauth: true
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -9,7 +9,10 @@ from diffusers.utils import numpy_to_pil
 from diffusers import WuerstchenDecoderPipeline, WuerstchenPriorPipeline
 from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
 from previewer.modules import Previewer
-os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from gallery_history import fetch_gallery_history, show_gallery_history
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 DESCRIPTION = "# Würstchen"
 DESCRIPTION += "\n<p style=\"text-align: center\"><a href='https://huggingface.co/warp-ai/wuerstchen' target='_blank'>Würstchen</a> is a new fast and efficient high resolution text-to-image architecture and model</p>"
@@ -26,8 +29,12 @@ PREVIEW_IMAGES = True
 dtype = torch.float16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    prior_pipeline = WuerstchenPriorPipeline.from_pretrained("warp-ai/wuerstchen-prior", torch_dtype=dtype)
-    decoder_pipeline = WuerstchenDecoderPipeline.from_pretrained("warp-ai/wuerstchen", torch_dtype=dtype)
+    prior_pipeline = WuerstchenPriorPipeline.from_pretrained(
+        "warp-ai/wuerstchen-prior", torch_dtype=dtype
+    )
+    decoder_pipeline = WuerstchenDecoderPipeline.from_pretrained(
+        "warp-ai/wuerstchen", torch_dtype=dtype
+    )
     if ENABLE_CPU_OFFLOAD:
         prior_pipeline.enable_model_cpu_offload()
         decoder_pipeline.enable_model_cpu_offload()
@@ -36,18 +43,27 @@ if torch.cuda.is_available():
         decoder_pipeline.to(device)
 
     if USE_TORCH_COMPILE:
-        prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
-        decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="reduce-overhead", fullgraph=True)
-    
+        prior_pipeline.prior = torch.compile(
+            prior_pipeline.prior, mode="reduce-overhead", fullgraph=True
+        )
+        decoder_pipeline.decoder = torch.compile(
+            decoder_pipeline.decoder, mode="reduce-overhead", fullgraph=True
+        )
+
     if PREVIEW_IMAGES:
         previewer = Previewer()
-        previewer.load_state_dict(torch.load("previewer/text2img_wurstchen_b_v1_previewer_100k.pt")["state_dict"])
+        previewer.load_state_dict(
+            torch.load("previewer/text2img_wurstchen_b_v1_previewer_100k.pt")[
+                "state_dict"
+            ]
+        )
         previewer.eval().requires_grad_(False).to(device).to(dtype)
 
         def callback_prior(i, t, latents):
             output = previewer(latents)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).cpu().numpy())
             return output
+
     else:
         previewer = None
         callback_prior = None
@@ -96,7 +112,7 @@ def generate(
             if isinstance(r, list):
                 yield r
         prior_output = r
-    
+
     decoder_output = decoder_pipeline(
         image_embeddings=prior_output.image_embeddings,
         prompt=prompt,
@@ -209,19 +225,21 @@ with gr.Blocks(css="style.css") as demo:
         cache_examples=CACHE_EXAMPLES,
     )
 
+    history = show_gallery_history()
+
     inputs = [
-            prompt,
-            negative_prompt,
-            seed,
-            width,
-            height,
-            prior_num_inference_steps,
-            # prior_timesteps,
-            prior_guidance_scale,
-            decoder_num_inference_steps,
-            # decoder_timesteps,
-            decoder_guidance_scale,
-            num_images_per_prompt,
+        prompt,
+        negative_prompt,
+        seed,
+        width,
+        height,
+        prior_num_inference_steps,
+        # prior_timesteps,
+        prior_guidance_scale,
+        decoder_num_inference_steps,
+        # decoder_timesteps,
+        decoder_guidance_scale,
+        num_images_per_prompt,
     ]
     prompt.submit(
         fn=randomize_seed_fn,
@@ -234,6 +252,8 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name="run",
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
     negative_prompt.submit(
         fn=randomize_seed_fn,
@@ -246,6 +266,8 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name=False,
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
     run_button.click(
         fn=randomize_seed_fn,
@@ -258,7 +280,9 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name=False,
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()

gallery_history.py

@@ -0,0 +1,121 @@
+"""
+How to use:
+1. Create a Space with a Persistent Storage attached. Filesystem will be available under `/data`.
+2. Add `hf_oauth: true` to the Space metadata (README.md). Make sure to have Gradio>=3.41.0 configured.
+3. Add `HISTORY_FOLDER` as a Space variable (example. `"/data/history"`).
+4. Add `filelock` as dependency in `requirements.txt`.
+5. Add history gallery to your Gradio app:
+    a. Add imports: `from gallery_history import fetch_gallery_history, show_gallery_history`
+    a. Add `history = show_gallery_history()` within `gr.Blocks` context.
+    b. Add `.then(fn=fetch_gallery_history, inputs=[prompt, result], outputs=history)` on the generate event.
+"""
+import json
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from uuid import uuid4
+
+import gradio as gr
+from filelock import FileLock
+
+_folder = os.environ.get("HISTORY_FOLDER")
+if _folder is None:
+    print(
+        "'HISTORY_FOLDER' environment variable not set. User history will be saved "
+        "locally and will be lost when the Space instance is restarted."
+    )
+    _folder = Path(__file__).parent / "history"
+HISTORY_FOLDER_PATH = Path(_folder)
+
+IMAGES_FOLDER_PATH = HISTORY_FOLDER_PATH / "images"
+IMAGES_FOLDER_PATH.mkdir(parents=True, exist_ok=True)
+
+
+def show_gallery_history():
+    gr.Markdown("## Past images\n\nYou must be logged in to activate it.")
+    with gr.Column():
+        with gr.Row():
+            gr.LoginButton()
+            gr.LogoutButton()
+        gallery = gr.Gallery(
+            label="Past images",
+            show_label=True,
+            elem_id="gallery",
+            object_fit="contain",
+            height="auto",
+            preview=True,
+            show_share_button=True,
+            show_download_button=True,
+        )
+        gallery.attach_load_event(fetch_gallery_history, every=None)
+        return gallery
+
+
+def fetch_gallery_history(
+    prompt: Optional[str] = None,
+    result: Optional[Dict] = None,
+    user: Optional[gr.OAuthProfile] = None,
+):
+    if user is None:
+        return []
+    try:
+        if prompt is not None and result is not None:  # None values means no new images
+            return _update_user_history(
+                user["preferred_username"], [(item["name"], prompt) for item in result]
+            )
+        else:
+            return _read_user_history(user["preferred_username"])
+    except Exception as e:
+        raise gr.Error(f"Error while fetching history: {e}") from e
+
+
+####################
+# Internal helpers #
+####################
+
+
+def _read_user_history(username: str) -> List[Tuple[str, str]]:
+    """Return saved history for that user."""
+    with _user_lock(username):
+        path = _user_history_path(username)
+        if path.exists():
+            return json.loads(path.read_text())
+        return []  # No history yet
+
+
+def _update_user_history(
+    username: str, new_images: List[Tuple[str, str]]
+) -> List[Tuple[str, str]]:
+    """Update history for that user and return it."""
+    with _user_lock(username):
+        # Read existing
+        path = _user_history_path(username)
+        if path.exists():
+            images = json.loads(path.read_text())
+        else:
+            images = []  # No history yet
+
+        # Copy images to persistent folder
+        for src_path, prompt in new_images:
+            images.append((_copy_image(src_path), prompt))
+
+        # Save and return
+        path.write_text(json.dumps(images))
+        return images
+
+
+def _user_history_path(username: str) -> Path:
+    return HISTORY_FOLDER_PATH / f"{username}.json"
+
+
+def _user_lock(username: str) -> FileLock:
+    """Ensure history is not corrupted if concurrent calls."""
+    return FileLock(f"{_user_history_path(username)}.lock")
+
+
+def _copy_image(src: str) -> str:
+    """Copy image to the persistent storage."""
+    dst = IMAGES_FOLDER_PATH / f"{uuid4().hex}_{Path(src).name}"  # keep file ext
+    shutil.copyfile(src, dst)
+    return str(dst)

requirements.txt

@@ -5,4 +5,5 @@ invisible-watermark==0.2.0
 Pillow==10.0.0
 torch==2.0.1
 transformers==4.32.1
-compel
+compel
+filelock