SCHEDULERS dropdown added

AVAILABLE_SCHEDULERS = {
"DDIM": DDIMScheduler,
"DDPM": DDPMScheduler,
"PNDM": PNDMScheduler,
"LMS Discrete": LMSDiscreteScheduler,
"Euler Discrete": EulerDiscreteScheduler,
"Euler Ancestral Discrete": EulerAncestralDiscreteScheduler,
"DPM Solver Multistep": DPMSolverMultistepScheduler,
"DPM Solver Singlestep": DPMSolverSinglestepScheduler,
}

app.py

@@ -9,6 +9,9 @@ import io
 import tempfile
 import zipfile
 import PIL
+import subprocess
+from huggingface_hub import Repository
+from utils import save_to_hub, save_to_local
 from dataclasses import dataclass
 from io import BytesIO
 def sanitize_filename(filename):
@@ -16,14 +19,27 @@ def sanitize_filename(filename):
     return re.sub(r'[\\/*?:"<>|]', "_", filename)
 
 from typing import Optional, Literal, Union
-from diffusers import DiffusionPipeline, DDIMScheduler
+from diffusers import (DiffusionPipeline, DDIMScheduler, DDPMScheduler, PNDMScheduler, 
+                       LMSDiscreteScheduler, EulerDiscreteScheduler, 
+                       EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, 
+                       DPMSolverSinglestepScheduler)
+
+AVAILABLE_SCHEDULERS = {
+    "DDIM": DDIMScheduler,
+    "DDPM": DDPMScheduler,
+    "PNDM": PNDMScheduler,
+    "LMS Discrete": LMSDiscreteScheduler,
+    "Euler Discrete": EulerDiscreteScheduler,
+    "Euler Ancestral Discrete": EulerAncestralDiscreteScheduler,
+    "DPM Solver Multistep": DPMSolverMultistepScheduler,
+    "DPM Solver Singlestep": DPMSolverSinglestepScheduler,
+}
 HF_TOKEN = os.environ.get("HF_TOKEN")
 import streamlit as st
 st.set_page_config(layout="wide")
 import torch
 from diffusers import (
     StableDiffusionPipeline,
-    EulerDiscreteScheduler,
     StableDiffusionInpaintPipeline,
     StableDiffusionImg2ImgPipeline,
 )
@@ -44,6 +60,7 @@ from st_clickable_images import clickable_images
 
 import streamlit.components.v1 as components
 
+
 prefix = 'image_generation'
 
 def dict_to_style(d):
@@ -78,7 +95,7 @@ def display_and_download_images(output_images, metadata):
             
 PIPELINE_NAMES = Literal["txt2img", "inpaint", "img2img"]
 
-DEFAULT_PROMPT = "a sprinkled donut sitting on top of a purple cherry apple with ice cubes, colorful hyperrealism, digital explosion of vibrant colors and abstract digital elements"
+DEFAULT_PROMPT = "sprinkled donut sitting on top of a purple cherry apple, colorful hyperrealism"
 DEFAULT_WIDTH, DEFAULT_HEIGHT = 512, 512
 OUTPUT_IMAGE_KEY = "output_img"
 LOADED_IMAGE_KEY = "loaded_image"
@@ -97,7 +114,8 @@ def set_image(key: str, img: Image.Image):
 
 @st.cache_resource(max_entries=1)
 def get_pipeline(
-    name: PIPELINE_NAMES,
+    name: str,
+    scheduler_name: str = None,
 ) -> Union[
     StableDiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
@@ -107,10 +125,18 @@ def get_pipeline(
         model_id = "FFusion/FFusion-BaSE"
         
         pipeline = DiffusionPipeline.from_pretrained(model_id)
-        # switch the scheduler in the pipeline to use the DDIMScheduler
-        pipeline.scheduler = DDIMScheduler.from_config(
-            pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
-        )
+
+        # Use specified scheduler if provided, else use DDIMScheduler
+        if scheduler_name:
+            SchedulerClass = AVAILABLE_SCHEDULERS[scheduler_name]
+            pipeline.scheduler = SchedulerClass.from_config(
+                pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
+            )
+        else:
+            pipeline.scheduler = DDIMScheduler.from_config(
+                pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
+            )
+            
         pipeline = pipeline.to("cuda")
         return pipeline
 
@@ -142,7 +168,7 @@ def generate(
 
     if enable_xformers:
         pipe.enable_xformers_memory_efficient_attention()
-        
+
     kwargs = dict(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -151,7 +177,6 @@ def generate(
         guidance_scale=guidance_scale,
         guidance_rescale=0.7
     )
-    print("kwargs", kwargs)
 
     if pipeline_name == "txt2img":
         kwargs.update(width=width, height=height)
@@ -164,6 +189,13 @@ def generate(
             f"Cannot generate image for pipeline {pipeline_name} and {prompt}"
         )
 
+    # Save images to Hugging Face Hub or locally
+    current_datetime = datetime.now()
+    metadata = {
+        "prompt": prompt,
+        "timestamp": str(current_datetime),
+    }
+
     output_images = []  # list to hold output image objects
     for _ in range(num_images):  # loop over number of images
         result = pipe(**kwargs)  # generate one image at a time
@@ -177,6 +209,10 @@ def generate(
             image.save(f"{filename}.png")
             output_images.append(image)  # add the image object to the list
 
+            # Save image to Hugging Face Hub
+            output_path = f"images/{i}.png"
+            save_to_hub(image, current_datetime, metadata, output_path)
+
     for image in output_images:
         with open(f"{filename}.txt", "w") as f:
             f.write(prompt)
@@ -185,6 +221,9 @@ def generate(
 
 
 
+
+
+
 def prompt_and_generate_button(prefix, pipeline_name: PIPELINE_NAMES, **kwargs):
     prompt = st.text_area(
         "Prompt",
@@ -203,9 +242,21 @@ def prompt_and_generate_button(prefix, pipeline_name: PIPELINE_NAMES, **kwargs):
         guidance_scale = st.slider(
             "Guidance scale", min_value=0.0, max_value=20.0, value=7.5, step=0.5, key=f"{prefix}-guidance-scale"
         )
+    # Add a select box for the schedulers
+    scheduler_name = st.selectbox(
+        "Choose a Scheduler",
+        options=list(AVAILABLE_SCHEDULERS.keys()),
+        index=0,  # Default index
+        key=f"{prefix}-scheduler",
+    )
+    scheduler_class = AVAILABLE_SCHEDULERS[scheduler_name]  # Get the selected scheduler class
+
+
+    pipe = get_pipeline(pipeline_name, scheduler_name=scheduler_name)        
+    
    # enable_attention_slicing = st.checkbox('Enable attention slicing (enables higher resolutions but is slower)', key=f"{prefix}-attention-slicing", value=True)
    # enable_xformers = st.checkbox('Enable xformers library (better memory usage)', key=f"{prefix}-xformers", value=True)
-        num_images = st.slider("Number of images to generate", min_value=1, max_value=4, value=1, key=f"{prefix}-num-images")
+    num_images = st.slider("Number of images to generate", min_value=1, max_value=4, value=1, key=f"{prefix}-num-images")
 
     images = []
 

utils.py

@@ -0,0 +1,227 @@
+import base64
+import gc
+import io
+import os
+import tempfile
+import zipfile
+from datetime import datetime
+from threading import Thread
+from huggingface_hub import Repository
+import subprocess
+
+import requests
+import streamlit as st
+import torch
+from huggingface_hub import HfApi
+from huggingface_hub.utils._errors import RepositoryNotFoundError
+from huggingface_hub.utils._validators import HFValidationError
+from loguru import logger
+from PIL.PngImagePlugin import PngInfo
+from st_clickable_images import clickable_images
+
+
+no_safety_checker = None
+
+
+CODE_OF_CONDUCT = """
+## Code of conduct
+The app should not be used to intentionally create or disseminate images that create hostile or alienating environments for people. This includes generating images that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.
+
+Using the app to generate content that is cruel to individuals is a misuse of this app. One shall not use this app to generate content that is intended to be cruel to individuals, or to generate content that is intended to be cruel to individuals in a way that is not obvious to the viewer.
+This includes, but is not limited to:
+- Generating demeaning, dehumanizing, or otherwise harmful representations of people or their environments, cultures, religions, etc.
+- Intentionally promoting or propagating discriminatory content or harmful stereotypes.
+- Impersonating individuals without their consent.
+- Sexual content without consent of the people who might see it.
+- Mis- and disinformation
+- Representations of egregious violence and gore
+- Sharing of copyrighted or licensed material in violation of its terms of use.
+- Sharing content that is an alteration of copyrighted or licensed material in violation of its terms of use.
+
+By using this app, you agree to the above code of conduct.
+
+"""
+
+
+def use_auth_token():
+    token_path = os.path.join(os.path.expanduser("~"), ".huggingface", "token")
+    if os.path.exists(token_path):
+        return True
+    if "HF_TOKEN" in os.environ:
+        return os.environ["HF_TOKEN"]
+    return False
+
+
+
+def download_file(file_url):
+    r = requests.get(file_url, stream=True)
+    with tempfile.NamedTemporaryFile(delete=False) as tmp:
+        for chunk in r.iter_content(chunk_size=1024):
+            if chunk:  # filter out keep-alive new chunks
+                tmp.write(chunk)
+    return tmp.name
+
+
+def cache_folder():
+    _cache_folder = os.path.join(os.path.expanduser("~"), ".ffusion")
+    os.makedirs(_cache_folder, exist_ok=True)
+    return _cache_folder
+
+
+def clear_memory(preserve):
+    torch.cuda.empty_cache()
+    gc.collect()
+    to_clear = ["inpainting", "text2img", "img2text"]
+    for key in to_clear:
+        if key not in preserve and key in st.session_state:
+            del st.session_state[key]
+
+
+
+    import subprocess
+
+from huggingface_hub import Repository
+
+def save_to_hub(image, current_datetime, metadata, output_path):
+    """Saves an image to Hugging Face Hub"""
+    try:
+        # Convert image to byte array
+        byte_arr = io.BytesIO()
+        
+        # Check if the image has metadata
+        if image.info:
+            # Save as PNG
+            image.save(byte_arr, format='PNG')
+        else:
+            # Save as JPG
+            image.save(byte_arr, format='JPEG')
+        
+        byte_arr = byte_arr.getvalue()
+
+        # Create a repository object
+        token = os.getenv("HF_TOKEN")
+        api = HfApi()
+
+        username = "FFusion"
+        repo_name = "FF"
+        try:
+            repo = Repository(f"{username}/{repo_name}", clone_from=f"{username}/{repo_name}", use_auth_token=token, repo_type="dataset")
+        except RepositoryNotFoundError:
+            repo = Repository(f"{username}/{repo_name}", clone_from=f"{username}/{repo_name}", use_auth_token=token, repo_type="dataset")
+
+        # Create the directory if it does not exist
+        os.makedirs(os.path.dirname(f"{repo.local_dir}/{output_path}"), exist_ok=True)
+
+        # Write image to repository
+        with open(f"{repo.local_dir}/{output_path}", "wb") as f:
+            f.write(byte_arr)
+
+        # Set Git username and email
+        subprocess.run(["git", "config", "user.name", "idle stoev"], check=True, cwd=repo.local_dir)
+        subprocess.run(["git", "config", "user.email", "[email protected]"], check=True, cwd=repo.local_dir)
+
+        # Commit and push changes
+        repo.git_add(pattern=".")
+        repo.git_commit(f"Add image at {current_datetime}")
+        print(f"Pushing changes to {username}/{repo_name}...")
+        repo.git_push()
+        print(f"Image saved to {username}/{repo_name}/{output_path}")
+
+    except Exception as e:
+        print(f"Failed to save image to Hugging Face Hub: {e}")
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+    
+
+
+def save_to_local(images, module, current_datetime, metadata, output_path):
+    _metadata = PngInfo()
+    _metadata.add_text("text2img", metadata)
+    os.makedirs(output_path, exist_ok=True)
+    os.makedirs(f"{output_path}/{module}", exist_ok=True)
+    os.makedirs(f"{output_path}/{module}/{current_datetime}", exist_ok=True)
+
+    for i, img in enumerate(images):
+        img.save(
+            f"{output_path}/{module}/{current_datetime}/{i}.png",
+            pnginfo=_metadata,
+        )
+
+    # save metadata as text file
+    with open(f"{output_path}/{module}/{current_datetime}/metadata.txt", "w") as f:
+        f.write(metadata)
+    logger.info(f"Saved images to {output_path}/{module}/{current_datetime}")
+
+
+def save_images(images, module, metadata, output_path):
+    if output_path is None:
+        logger.warning("No output path specified, skipping saving images")
+        return
+
+    api = HfApi()
+    dset_info = None
+    try:
+        dset_info = api.dataset_info(output_path)
+    except (HFValidationError, RepositoryNotFoundError):
+        logger.warning("No valid hugging face repo. Saving locally...")
+
+    current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+    if not dset_info:
+        save_to_local(images, module, current_datetime, metadata, output_path)
+    else:
+        Thread(target=save_to_hub, args=(api, images, module, current_datetime, metadata, output_path)).start()
+
+
+def display_and_download_images(output_images, metadata, download_col=None):
+    # st.image(output_images, width=128, output_format="PNG")
+
+    with st.spinner("Preparing images for download..."):
+        # save images to a temporary directory
+        with tempfile.TemporaryDirectory() as tmpdir:
+            gallery_images = []
+            for i, image in enumerate(output_images):
+                image.save(os.path.join(tmpdir, f"{i + 1}.png"), pnginfo=metadata)
+                with open(os.path.join(tmpdir, f"{i + 1}.png"), "rb") as img:
+                    encoded = base64.b64encode(img.read()).decode()
+                    gallery_images.append(f"data:image/jpeg;base64,{encoded}")
+
+            # zip the images
+            zip_path = os.path.join(tmpdir, "images.zip")
+            with zipfile.ZipFile(zip_path, "w") as zip:
+                for filename in os.listdir(tmpdir):
+                    if filename.endswith(".png"):
+                        zip.write(os.path.join(tmpdir, filename), filename)
+
+            # convert zip to base64
+            with open(zip_path, "rb") as f:
+                encoded = base64.b64encode(f.read()).decode()
+
+            _ = clickable_images(
+                gallery_images,
+                titles=[f"Image #{str(i)}" for i in range(len(gallery_images))],
+                div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
+                img_style={"margin": "5px", "height": "200px"},
+            )
+
+            # add download link
+            st.markdown(
+                f"""
+                <a href="data:application/zip;base64,{encoded}" download="images.zip">
+                    Download Images
+                </a>
+                """,
+                unsafe_allow_html=True,
+            )