Update app.py

app.py

@@ -1,127 +1,113 @@
-import os
-import sys
-import random
 import torch
-import numpy as np
-from PIL import Image
+from diffusers import FluxPipeline
 import gradio as gr
+import random
+import numpy as np
+import os
+import spaces
 
-# Check and add the ComfyUI repository path to sys.path
-repo_path = './ComfyUI/totoro_extras'
-print(f"Checking for repository path: {repo_path}")
-if not os.path.exists(repo_path):
-    raise FileNotFoundError(f"Repository path '{repo_path}' not found. Make sure the ComfyUI repository is cloned correctly.")
-sys.path.append(repo_path)
-print(f"Repository path added to sys.path: {repo_path}")
+# Check for GPU availability
+if torch.cuda.is_available():
+    device = "cuda"
+    print("Using GPU")
+else:
+    device = "cpu"
+    print("Using CPU")
 
-# Import nodes and custom modules
-from nodes import NODE_CLASS_MAPPINGS
-from totoro_extras import nodes_custom_sampler, nodes_flux
+# Set environment variables
+HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure to set this in your environment
 
-# Initialize necessary components from the nodes
-CheckpointLoaderSimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
-LoraLoader = NODE_CLASS_MAPPINGS["LoraLoader"]()
-FluxGuidance = nodes_flux.NODE_CLASS_MAPPINGS["FluxGuidance"]()
-RandomNoise = nodes_custom_sampler.NODE_CLASS_MAPPINGS["RandomNoise"]()
-BasicGuider = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicGuider"]()
-KSamplerSelect = nodes_custom_sampler.NODE_CLASS_MAPPINGS["KSamplerSelect"]()
-BasicScheduler = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicScheduler"]()
-SamplerCustomAdvanced = nodes_custom_sampler.NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
-VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
-VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
-EmptyLatentImage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
+MAX_SEED = np.iinfo(np.int32).max
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
 
-# Load checkpoints and models
-with torch.inference_mode():
-    checkpoint_path = "models/checkpoints/flux1-dev-fp8-all-in-one.safetensors"
-    unet, clip, vae = CheckpointLoaderSimple.load_checkpoint(checkpoint_path)
+# Initialize the pipeline and download the model
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.to(device)
 
-def closestNumber(n, m):
-    q = int(n / m)
-    n1 = m * q
-    if (n * m) > 0:
-        n2 = m * (q + 1)
-    else:
-        n2 = m * (q - 1)
-    if abs(n - n1) < abs(n - n2):
-        return n1
-    return n2
+# Enable memory optimizations
+pipe.enable_attention_slicing()
 
-@torch.inference_mode()
-def generate(positive_prompt, width, height, seed, steps, sampler_name, scheduler, guidance, lora_strength_model, lora_strength_clip):
-    global unet, clip
+# Define the image generation function
+@spaces.GPU(duration=180)
+def generate_image(prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt, progress=gr.Progress(track_tqdm=True)):
     if seed == 0:
-        seed = random.randint(0, 18446744073709551615)
-    print(f"Seed used: {seed}")
-
-    # Load LoRA models
-    lora_path = "models/loras/flux_realism_lora.safetensors"
-    unet_lora, clip_lora = LoraLoader.load_lora(unet, clip, lora_path, lora_strength_model, lora_strength_clip)
-
-    # Encode the prompt
-    cond, pooled = clip_lora.encode_from_tokens(clip_lora.tokenize(positive_prompt), return_pooled=True)
-    cond = [[cond, {"pooled_output": pooled}]]
-    cond = FluxGuidance.append(cond, guidance)[0]
-
-    # Generate noise
-    noise = RandomNoise.get_noise(seed)[0]
-
-    # Get guider and sampler
-    guider = BasicGuider.get_guider(unet_lora, cond)[0]
-    sampler = KSamplerSelect.get_sampler(sampler_name)[0]
-
-    # Get scheduling sigmas
-    sigmas = BasicScheduler.get_sigmas(unet_lora, scheduler, steps, 1.0)[0]
-
-    # Generate latent image
-    latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0]
-
-    # Sample and decode the image
-    sample, sample_denoised = SamplerCustomAdvanced.sample(noise, guider, sampler, sigmas, latent_image)
-    decoded = VAEDecode.decode(vae, sample)[0].detach()
-
-    # Convert to image and return
-    return Image.fromarray(np.array(decoded * 255, dtype=np.uint8)[0])
-
-# Define Gradio interface
-with gr.Blocks(analytics_enabled=False) as demo:
+        seed = random.randint(1, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+    
+    with torch.inference_mode():
+        output = pipe(
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            height=height,
+            width=width,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            num_images_per_prompt=num_images_per_prompt
+        ).images
+    
+    return output
+
+# Create the Gradio interface
+examples = [
+    ["A cat holding a sign that says hello world"],
+    ["a tiny astronaut hatching from an egg on the moon"],
+    ["An astronaut on Mars in a futuristic cyborg suit."],
+]
+
+css = '''
+.gradio-container{max-width: 1000px !important}
+h1{text-align:center}
+'''
+
+with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column():
-            positive_prompt = gr.Textbox(
-                lines=3, 
-                interactive=True, 
-                value="cute anime girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black dress with a gold leaf pattern and a white apron eating a slice of an apple pie in the kitchen of an old dark victorian mansion with a bright window and very expensive stuff everywhere", 
-                label="Prompt"
+            gr.HTML(
+                """
+                <h1 style='text-align: center'>
+                FLUX.1-dev
+                </h1>
+                """
             )
-            width = gr.Slider(minimum=256, maximum=2048, value=1024, step=16, label="width")
-            height = gr.Slider(minimum=256, maximum=2048, value=1024, step=16, label="height")
-            seed = gr.Slider(minimum=0, maximum=18446744073709551615, value=0, step=1, label="seed (0=random)")
-            steps = gr.Slider(minimum=4, maximum=50, value=20, step=1, label="steps")
-            guidance = gr.Slider(minimum=0, maximum=20, value=3.5, step=0.5, label="guidance")
-            lora_strength_model = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.1, label="lora_strength_model")
-            lora_strength_clip = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.1, label="lora_strength_clip")
-            sampler_name = gr.Dropdown(
-                ["euler", "heun", "heunpp2", "dpm_2", "lms", "dpmpp_2m", "ipndm", "deis", "ddim", "uni_pc", "uni_pc_bh2"], 
-                label="sampler_name", 
-                value="euler"
+            gr.HTML(
+                """
+                Made by <a href='https://linktr.ee/Nick088' target='_blank'>Nick088</a>
+                <br> <a href="https://discord.gg/osai"> <img src="https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge" alt="Discord"> </a>
+                """
             )
-            scheduler = gr.Dropdown(
-                ["normal", "sgm_uniform", "simple", "ddim_uniform"], 
-                label="scheduler", 
-                value="simple"
-            )
-            generate_button = gr.Button("Generate")
+    with gr.Group():
         with gr.Column():
-            output_image = gr.Image(label="Generated image", interactive=False)
+            prompt = gr.Textbox(label="Prompt", info="Describe the image you want", placeholder="A cat...")
+            run_button = gr.Button("Run")
+        result = gr.Gallery(label="Generated AI Images", elem_id="gallery")
+    with gr.Accordion("Advanced options", open=False):
+        with gr.Row():
+            num_inference_steps = gr.Slider(label="Number of Inference Steps", info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference", minimum=1, maximum=50, value=25, step=1)
+            guidance_scale = gr.Slider(label="Guidance Scale", info="Controls how much the image generation process follows the text prompt. Higher values make the image stick more closely to the input text.", minimum=0.0, maximum=7.0, value=3.5, step=0.1)
+        with gr.Row():
+            width = gr.Slider(label="Width", info="Width of the Image", minimum=256, maximum=1024, step=32, value=1024)
+            height = gr.Slider(label="Height", info="Height of the Image", minimum=256, maximum=1024, step=32, value=1024)
+        with gr.Row():
+            seed = gr.Slider(value=42, minimum=0, maximum=MAX_SEED, step=1, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")
+            num_images_per_prompt = gr.Slider(label="Images Per Prompt", info="Number of Images to generate with the settings", minimum=1, maximum=4, step=1, value=2)
+
+    gr.Examples(
+        examples=examples,
+        fn=generate_image,
+        inputs=[prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt],
+        outputs=[result],
+        cache_examples=CACHE_EXAMPLES
+    )
 
-    generate_button.click(
-        fn=generate, 
-        inputs=[
-            positive_prompt, width, height, seed, steps, 
-            sampler_name, scheduler, guidance, 
-            lora_strength_model, lora_strength_clip
-        ], 
-        outputs=output_image
+    gr.on(
+        triggers=[
+            prompt.submit,
+            run_button.click,
+        ],
+        fn=generate_image,
+        inputs=[prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt],
+        outputs=[result],
     )
 
-demo.queue().launch(inline=False, share=True, debug=True)
+demo.queue().launch(share=False)