Update app.py

app.py

@@ -1,51 +1,36 @@
-import os
 import gradio as gr
-import numpy as np
 import json
-from accelerate import dispatch_model, infer_auto_device_map
-from accelerate.utils import get_balanced_memory
-from torch.cuda.amp import autocast
+import logging
+import argparse
 import torch
-import spaces  # Import this first to avoid CUDA initialization issues
+import os
+from os import path
+from PIL import Image
+import spaces
+import copy
 import random
 import time
-from PIL import Image
-from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
-from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from huggingface_hub import hf_hub_download
+from diffusers import FluxTransformer2DModel, FluxPipeline
+import safetensors.torch
+from safetensors.torch import load_file
+import gc
 
-# Use the 'waffles' environment variable as the access token
-hf_token = os.getenv('waffles')
+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
 
-# Ensure the token is loaded correctly
-if not hf_token:
-    raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.")
+torch.backends.cuda.matmul.allow_tf32 = True
 
-# Define the device
-dtype = torch.bfloat16
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    n_gpu = torch.cuda.device_count()
-    torch.cuda.get_device_name(0)
-else:
-    device = torch.device("cpu")
-
-count0 = torch.zeros(1).to(device)
-count1 = torch.zeros(1).to(device)
-count2 = torch.zeros(1).to(device)
+pipe = FluxPipeline.from_pretrained("John6666/nsfw-master-flux-lora-merged-with-flux1-dev-fp16-v10-fp8-flux", torch_dtype=torch.bfloat16)
+pipe.to(device="cuda", dtype=torch.bfloat16)
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
-
-# Initialize the base model with authentication and specify the device
-# Initialize the base model with authentication and specify the device
-pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token).to(device)
-
-MAX_SEED = 2**32 - 1
-MAX_IMAGE_SIZE = 2048
-
+    
+MAX_SEED = 2**32-1
 
 class calculateDuration:
     def __init__(self, activity_name=""):
@@ -63,36 +48,62 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
-@spaces.GPU(duration=90)
-def generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress):
-    generator = torch.Generator(device=device).manual_seed(seed)
-    images = []
+
+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 images"):
-        for _ in range(num_images):
-            # Generate each 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]
-            images.append(image)
-    return images
-
-def run_lora(prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, width, height, lora_scale, num_images, progress=gr.Progress(track_tqdm=True)):
-    if not selected_repo:
+    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 = next((lora for lora in loras if lora["repo"] == selected_repo), None)
-    if not selected_lora:
-        raise gr.Error("Selected LoRA not found.")
-
+    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']}"):
@@ -106,15 +117,10 @@ def run_lora(prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, widt
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
     
-    images = generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress)
-    pipe.to("cuda")
+    image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    pipe.to("cpu")
     pipe.unload_lora_weights()
-    return images, seed  
-
-def update_selection(evt: gr.SelectData):
-    index = evt.index
-    selected_lora = loras[index]
-    return f"Selected LoRA: {selected_lora['title']}", selected_lora["repo"]
+    return image, seed  
 
 run_lora.zerogpu = True
 
@@ -123,78 +129,70 @@ css = '''
 #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: auto; width: auto;}
-#gallery .gallery-item{width: 50px; height: 50px; margin: 0px;} /* Make buttons 50% height and width */
-#gallery img{width: 100%; height: 100%; object-fit: cover;} /* Resize images to fit buttons */
-#info_blob {
-    background-color: #f0f0f0;
-    border: 2px solid #ccc;
-    padding: 10px;
-    margin: 10px 0;
-    text-align: center;
-    font-size: 1.2em;
-    font-weight: bold;
-    color: #333;
-    border-radius: 8px;
-}
+#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 on Schnell LoRas </h1>""",
+        """<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 stating what the app is running
     info_blob = gr.HTML(
-        """<div id="info_blob"> Activist, Futurist, and Realist LoRa-stocked Quick-Use Image Manufactory (over Flux Schnell)</div>"""
+        """<div id="info_blob"> Activist & Futurealist LoRa-stocked Img Manufactory (on Flux Dev HYPER (8-Step))</div>"""
     )
 
-    selected_lora_text = gr.Markdown("Selected LoRA: None")
-    selected_repo = gr.State(value="")
-
-    # Prompt takes the full line
-    prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Type a prompt after selecting a LoRA", elem_id="full_line_prompt")
-
+        # 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=1):  # LoRA collection on the left
+        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 Gallery",
+                label="LoRA Inventory",
                 allow_preview=False,
                 columns=3,
                 elem_id="gallery"
             )
-        with gr.Column(scale=1):  # Generated images on the right
-            result = gr.Gallery(label="Generated Images")
-            seed = gr.Number(label="Seed", value=0, interactive=False)
-
-    with gr.Column():
-        with gr.Row():
-            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=1)
-            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=4)
-        
-        with gr.Row():
-            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
-            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)
-            lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.95)
-            num_images = gr.Slider(label="Number of Images", minimum=1, maximum=4, step=1, value=1)
+            
+        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=3.5)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=6)
+                
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
+                    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=1.5, step=0.01, value=0.9)
 
     gallery.select(
-        fn=update_selection,
-        inputs=[],
-        outputs=[selected_lora_text, selected_repo]
+        update_selection,
+        inputs=[width, height],
+        outputs=[prompt, selected_info, selected_index, width, height]
     )
 
-    generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
-    generate_button.click(
-        run_lora,
-        inputs=[prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, width, height, lora_scale, num_images],
+    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()
-app.launch()
+app.queue(default_concurrency_limit=2).launch(show_error=True)
+app.launch()