Create app.py

app.py

@@ -0,0 +1,213 @@
+import gradio as gr
+import json
+import logging
+import argparse
+import torch
+import os
+from os import path
+from PIL import Image
+import numpy as np
+import spaces
+import copy
+import random
+import time
+from typing import Any, Dict, List, Optional, Union
+from huggingface_hub import hf_hub_download
+from diffusers import DiffusionPipeline, FluxTransformer2DModel, FluxPipeline, AutoencoderTiny
+import safetensors.torch
+from safetensors.torch import load_file
+from custom_pipeline import FluxWithCFGPipeline
+import gc
+
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+dtype = torch.float16
+pipe = FluxWithCFGPipeline.from_pretrained(
+    "ostris/OpenFLUX.1", torch_dtype=dtype
+)
+# pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
+# pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
+# pipe.set_adapters("fast")
+# pipe.fuse_lora(adapter_names=["fast"], lora_scale=1.0)
+pipe.to("cuda")
+ pipe.transformer.to(memory_format=torch.channels_last)
+ pipe.transformer = torch.compile(
+     pipe.transformer, mode="max-autotune", fullgraph=True
+ )
+torch.cuda.empty_cache()
+
+# Load LoRAs from JSON file
+with open('loras.json', 'r') as f:
+    loras = json.load(f)
+    
+MAX_SEED = 2**32-1
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
+
+def update_selection(evt: gr.SelectData, width, height):
+    selected_lora = loras[evt.index]
+    new_placeholder = f"Type a prompt for {selected_lora['title']}"
+    lora_repo = selected_lora["repo"]
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
+    if "aspect" in selected_lora:
+        if selected_lora["aspect"] == "portrait":
+            width = 768
+            height = 1024
+        elif selected_lora["aspect"] == "landscape":
+            width = 1024
+            height = 768
+    return (
+        gr.update(placeholder=new_placeholder),
+        updated_text,
+        evt.index,
+        width,
+        height,
+    )
+
+@spaces.GPU(duration=70)
+def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
+    pipe.to("cuda")
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
+    with calculateDuration("Generating image"):
+        # Generate image
+        image = pipe(
+            prompt=f"{prompt} {trigger_word}",
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
+    return image
+
+def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    if selected_index is None:
+        raise gr.Error("You must select a LoRA before proceeding.")
+
+    selected_lora = loras[selected_index]
+    lora_path = selected_lora["repo"]
+    trigger_word = selected_lora["trigger_word"]
+    if(trigger_word):
+        if "trigger_position" in selected_lora:
+            if selected_lora["trigger_position"] == "prepend":
+                prompt_mash = f"{trigger_word} {prompt}"
+            else:
+                prompt_mash = f"{prompt} {trigger_word}"
+        else:
+            prompt_mash = f"{trigger_word} {prompt}"
+    else:
+        prompt_mash = prompt
+
+    # Load LoRA weights
+    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
+        if "weights" in selected_lora:
+            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+        else:
+            pipe.load_lora_weights(lora_path)
+        
+    # Set random seed for reproducibility
+    with calculateDuration("Randomizing seed"):
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+    
+    image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    pipe.to("cpu")
+    pipe.unload_lora_weights()
+    return image, seed  
+
+run_lora.zerogpu = True
+
+css = '''
+#gen_btn{height: 100%}
+#title{text-align: center}
+#title h1{font-size: 3em; display:inline-flex; align-items:center}
+#title img{width: 100px; margin-right: 0.5em}
+#gallery .grid-wrap{height: 10vh}
+'''
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
+    title = gr.HTML(
+        """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> SOONfactory </h1>""",
+        elem_id="title",
+    )
+    	    # Info blob stating what the app is running
+    info_blob = gr.HTML(
+        """<div id="info_blob"> Activist & Futurealist LoRa-stocked Img Manufactory (currently on the Acorn is Spinning Schnell V1 model checkpoint, 4-step, by Seeker70) )</div>"""
+    )
+
+        # Info blob stating what the app is running
+    info_blob = gr.HTML(
+        """<div id="info_blob">Prephrase prompts w/: 1.RCA style 2. HST style autochrome 3. HST style 4.TOK hybrid 5.2004 photo 6.HST style 7.LEN Vladimir Lenin 8.TOK portra 9.HST portrait 10.flmft 11.HST in Peterhof 12.HST Soviet kodachrome 13. SOTS art 14.HST 15.photo 16.pficonics 17.wh3r3sw4ld0 18.retrofuturism 19-24.HST style photo 25.vintage cover </div>"""
+    )
+    selected_index = gr.State(None)
+    with gr.Row():
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Select LoRa/Style & type prompt!")
+        with gr.Column(scale=1, elem_id="gen_column"):
+            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
+    with gr.Row():
+        with gr.Column(scale=3):
+            selected_info = gr.Markdown("")
+            gallery = gr.Gallery(
+                [(item["image"], item["title"]) for item in loras],
+                label="LoRA Inventory",
+                allow_preview=False,
+                columns=3,
+                elem_id="gallery"
+            )
+            
+        with gr.Column(scale=4):
+            result = gr.Image(label="Generated Image")
+
+    with gr.Row():
+        with gr.Accordion("Advanced Settings", open=True):
+            with gr.Column():
+                with gr.Row():
+                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=1.0)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=5)
+                
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=768)
+                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
+                
+                with gr.Row():
+                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
+                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2.0, step=0.01, value=0.6)
+
+    gallery.select(
+        update_selection,
+        inputs=[width, height],
+        outputs=[prompt, selected_info, selected_index, width, height]
+    )
+
+    gr.on(
+        triggers=[generate_button.click, prompt.submit],
+        fn=run_lora,
+        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
+        outputs=[result, seed]
+    )
+
+app.queue(default_concurrency_limit=None).launch(show_error=True)
+app.launch()