Update app.py

app.py

@@ -1,206 +1,80 @@
 import os
 import spaces
 import torch
-import json
-import logging
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, StableDiffusionXLImg2ImgPipeline, AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 import gradio as gr
 import random
-from datetime import datetime
-from PIL import Image, PngImagePlugin
+import tqdm
 
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Enable TQDM progress tracking
+tqdm.monitor_interval = 0
 
-# Configuration
-OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
-MAX_SEED = 2**32 - 1
+#HF_TOKEN import
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-def seed_everything(seed):
-    if seed is None:
-        seed = random.randint(0, MAX_SEED)
-    torch.manual_seed(seed)
-    random.seed(seed)
-    return torch.Generator(device='cuda').manual_seed(seed)
-
-def save_image(image, metadata, output_dir, is_colab=False):
-    os.makedirs(output_dir, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"generated_{timestamp}.png"
-    filepath = os.path.join(output_dir, filename)
-    
-    # Save with metadata
-    png_info = PngImagePlugin.PngInfo()
-    png_info.add_text("parameters", json.dumps(metadata))
-    image.save(filepath, "PNG", pnginfo=png_info)
-    
-    return filepath
-
-# Load the diffusion pipeline with optimized VAE
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "kayfahaarukku/irAsu-1.0",
+# Load the diffusion pipeline
+pipe = StableDiffusionXLPipeline.from_single_file(
+    "https://huggingface.co/kayfahaarukku/AkashicPulse-v1.0/resolve/main/AkashicPulse-v1.0-ft-ft.safetensors",
     torch_dtype=torch.float16,
     custom_pipeline="lpw_stable_diffusion_xl",
+    use_safetensors=True,
+    use_auth_token=HF_TOKEN,
 )
-
-# Load optimized VAE
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix",
-    torch_dtype=torch.float16,
-)
-pipe.vae = vae
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-# Style presets
-styles = {
-    "(None)": ("", ""),
-    "Detailed": ("highly detailed, intricate details, ", ""),
-    "Simple": ("simple style, minimalistic, ", "complex, detailed"),
-    "Soft": ("soft lighting, dreamy atmosphere, ", "harsh lighting, sharp contrast"),
-}
-
-# Quality presets
-quality_presets = {
-    "Standard": (
-        "best quality, amazing quality, very aesthetic",
-        "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts"
-    ),
-    "High Detail": (
-        "masterpiece, best quality, amazing quality, very aesthetic, highly detailed",
-        "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality"
-    ),
-    "Basic": (
-        "good quality",
-        "nsfw, lowres, bad quality"
-    )
-}
-
 # Function to generate an image
-@spaces.GPU
-def generate_image(
-    prompt,
-    negative_prompt,
-    use_quality_preset,
-    resolution,
-    guidance_scale,
-    num_inference_steps,
-    seed,
-    randomize_seed,
-    style_preset="(None)",
-    use_upscaler=False,
-    upscaler_strength=0.55,
-    upscale_by=1.5,
-    progress=gr.Progress()
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+@spaces.GPU  # Adjust the duration as needed
+def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
+    pipe.to('cuda')  # Move the model to GPU when the function is called
     
-    # Apply style preset
-    style_prompt, style_negative = styles[style_preset]
-    prompt = f"{style_prompt}{prompt}"
-    negative_prompt = f"{negative_prompt}, {style_negative}" if style_negative else negative_prompt
-    
-    if use_quality_preset:
-        quality_prompt, quality_negative = quality_presets["Standard"]
-        prompt = f"{prompt}, {quality_prompt}"
-        negative_prompt = f"{negative_prompt}, {quality_negative}"
+    if randomize_seed:
+        seed = random.randint(0, 99999999)
+    if use_defaults:
+        prompt = f"{prompt}, best quality, amazing quality, very aesthetic"
+        negative_prompt = f"nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], {negative_prompt}"
+    generator = torch.manual_seed(seed)
     
-    generator = seed_everything(seed)
+    def callback(step, timestep, latents):
+        progress(step / num_inference_steps)
+        return
     
     width, height = map(int, resolution.split('x'))
-    
-    metadata = {
-        "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "resolution": f"{width} x {height}",
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "seed": seed,
-        "style_preset": style_preset,
-        "use_quality_preset": use_quality_preset
-    }
-    
-    try:
-        if use_upscaler:
-            # Initial generation
-            latents = pipe(
-                prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="latent"
-            ).images
-            
-            # Setup img2img pipeline for upscaling
-            upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-            
-            # Calculate new dimensions
-            new_width = int(width * upscale_by)
-            new_height = int(height * upscale_by)
-            
-            # Upscale
-            image = upscaler_pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=latents,
-                strength=upscaler_strength,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator
-            ).images[0]
-            
-            metadata["upscaler"] = {
-                "strength": upscaler_strength,
-                "scale_factor": upscale_by,
-                "final_resolution": f"{new_width}x{new_height}"
-            }
-        else:
-            image = pipe(
-                prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                callback=lambda step, timestep, latents: progress(step / num_inference_steps)
-            ).images[0]
-        
-        # Save image with metadata
-        image_path = save_image(image, metadata, OUTPUT_DIR)
-        logger.info(f"Image saved as {image_path} with metadata")
-        
-        return image, seed, json.dumps(metadata, indent=2)
-        
-    except Exception as e:
-        logger.exception(f"An error occurred: {e}")
-        raise
-    finally:
-        if use_upscaler:
-            del upscaler_pipe
-        torch.cuda.empty_cache()
+    image = pipe(
+        prompt, 
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height, 
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        callback=callback,
+        callback_steps=1
+    ).images[0]
+
+    torch.cuda.empty_cache()
+
+    metadata_text = f"{prompt}\nNegative prompt: {negative_prompt}\nSteps: {num_inference_steps}, Sampler: Euler a, Size: {width}x{height}, Seed: {seed}, CFG scale: {guidance_scale}"
+
+    return image, seed, metadata_text
 
 # Define Gradio interface
-with gr.Blocks(title="irAsu 1.0 Enhanced Demo", theme="NoCrypt/miku@1.2.1") as demo:
-    gr.HTML("<h1>irAsu 1.0 Enhanced Demo</h1>")
-    
+def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
+    image, seed, metadata_text = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
+    return image, seed, gr.update(value=metadata_text)
+
+def reset_inputs():
+    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='832x1216'), gr.update(value=4), gr.update(value=28), gr.update(value=0), gr.update(value=True), gr.update(value='')
+
+with gr.Blocks(title="irAsu 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
+    gr.HTML(
+        "<h1>irAsu 1.0 Demo</h1>"
+        "This demo is intended to showcase what the model is capable of and is not intended to be the main generation platform. Results produced with Diffusers are not the best, and it's highly recommended for you to get the model running inside Stable Diffusion WebUI or ComfyUI."
+        )
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
             negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
-            
-            with gr.Accordion("Style & Quality", open=True):
-                style_selector = gr.Radio(
-                    choices=list(styles.keys()),
-                    value="(None)",
-                    label="Style Preset"
-                )
-                use_quality_preset = gr.Checkbox(label="Use Quality Preset", value=True)
-            
+            use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
             resolution_input = gr.Radio(
                 choices=[
                     "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
@@ -209,65 +83,49 @@ with gr.Blocks(title="irAsu 1.0 Enhanced Demo", theme="NoCrypt/[email protected]") as d
                 label="Resolution",
                 value="832x1216"
             )
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=4)
-                num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
-                seed_input = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, label="Seed", value=0)
-                randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-                
-                use_upscaler_input = gr.Checkbox(label="Use Upscaler", value=False)
-                with gr.Group(visible=False) as upscaler_settings:
-                    upscaler_strength_input = gr.Slider(minimum=0, maximum=1, step=0.05, label="Upscaler Strength", value=0.55)
-                    upscale_by_input = gr.Slider(minimum=1, maximum=1.5, step=0.1, label="Upscale Factor", value=1.5)
-            
+            guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=4)
+            num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
+            seed_input = gr.Slider(minimum=0, maximum=999999999, step=1, label="Seed", value=0, interactive=True)
+            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
             generate_button = gr.Button("Generate")
             reset_button = gr.Button("Reset")
 
         with gr.Column():
             output_image = gr.Image(type="pil", label="Generated Image")
             with gr.Accordion("Parameters", open=False):
-                metadata_textbox = gr.Textbox(lines=6, label="Image Parameters", interactive=False)
+                gr.Markdown(
+                    """
+                    This parameter is compatible with Stable Diffusion WebUI's parameter importer.
+                    """
+                )
+                metadata_textbox = gr.Textbox(lines=6, label="Image Parameters", interactive=False, max_lines=6)
+            gr.Markdown(
+                """
+                ### Recommended prompt formatting:
+                `1girl/1boy, character name, from what series, everything else in any order, best quality, amazing quality, very aesthetic,`
 
-    # Handle upscaler visibility
-    use_upscaler_input.change(
-        fn=lambda x: gr.Group(visible=x),
-        inputs=[use_upscaler_input],
-        outputs=[upscaler_settings]
-    )
+                **PS:** `best quality, amazing quality, very aesthetic,` is automatically added when "Use Default Quality Tags and Negative Prompt" is enabled
+
+                ### Recommended settings:
+                - Steps: 25-30
+                - CFG: 3.5-5
+                - Sweet spot: 28 steps, 4 CFG
+                """
+            )
 
-    # Generate button click event
     generate_button.click(
-        generate_image,
+        interface_fn,
         inputs=[
-            prompt_input,
-            negative_prompt_input,
-            use_quality_preset,
-            resolution_input,
-            guidance_scale_input,
-            num_inference_steps_input,
-            seed_input,
-            randomize_seed_input,
-            style_selector,
-            use_upscaler_input,
-            upscaler_strength_input,
-            upscale_by_input
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
         ],
         outputs=[output_image, seed_input, metadata_textbox]
     )
-
-    # Reset button click event
+    
     reset_button.click(
-        lambda: (
-            "", "", True, "832x1216", 4, 28, 0, True,
-            "(None)", False, 0.55, 1.5, None
-        ),
+        reset_inputs,
+        inputs=[],
         outputs=[
-            prompt_input, negative_prompt_input, use_quality_preset,
-            resolution_input, guidance_scale_input, num_inference_steps_input,
-            seed_input, randomize_seed_input, style_selector,
-            use_upscaler_input, upscaler_strength_input, upscale_by_input,
-            metadata_textbox
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, metadata_textbox
         ]
     )