Create app.py

app.py

@@ -0,0 +1,103 @@
+import math
+import random
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import (
+    StableDiffusionXLImg2ImgPipeline, 
+    StableDiffusionXLPipeline, 
+    EDMEulerScheduler, 
+    AutoencoderKL, 
+    DPMSolverMultistepScheduler
+)
+from huggingface_hub import hf_hub_download, InferenceClient
+
+# Load models and pipelines
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+pipe = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V4.0", torch_dtype=torch.float16, vae=vae)
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
+pipe.to("cuda")
+
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True)
+refiner.to("cuda")
+
+pipe_fast = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V4.0_Lightning", torch_dtype=torch.float16, vae=vae, use_safetensors=True)
+pipe_fast.to("cuda")
+
+# Inference Client for prompt optimization
+client1 = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+system_instructions = "<|system|>\nOptimize the following prompt for better image generation. Make it concise, high-quality, and descriptive.\n<|user|>\n"
+
+# Function to optimize prompts
+def promptifier(prompt):
+    return client1.text_generation(f"{system_instructions}{prompt}\n<|assistant|>\n", max_new_tokens=100)
+
+# Image generation/editing function
+def king(task, input_image, instruction, negative_prompt="", enhance_prompt=True, steps=25, randomize_seed=True, seed=2404, width=1024, height=1024, guidance_scale=6, fast=True):
+    if randomize_seed:
+        seed = random.randint(0, 999999)
+    generator = torch.manual_seed(seed)
+    
+    if enhance_prompt:
+        instruction = promptifier(instruction)
+    
+    if task == "Image Editing" and input_image:
+        input_image = Image.open(input_image).convert('RGB')
+        output_image = pipe_edit(instruction, negative_prompt=negative_prompt, image=input_image, guidance_scale=guidance_scale, width=width, height=height, num_inference_steps=steps, generator=generator).images
+        refined_image = refiner(prompt=f"{instruction}, 4k, hd, high quality", negative_prompt=negative_prompt, guidance_scale=7.5, num_inference_steps=steps, image=output_image, generator=generator).images[0]
+        return seed, refined_image
+    else:
+        if fast:
+            output_image = pipe_fast(prompt=instruction, guidance_scale=guidance_scale/2, num_inference_steps=int(steps/2.5), width=width, height=height, generator=generator).images[0]
+        else:
+            latent_image = pipe_fast(prompt=instruction, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=steps, width=width, height=height, generator=generator, output_type="latent").images
+            refined_image = refiner(prompt=instruction, negative_prompt=negative_prompt, guidance_scale=7.5, num_inference_steps=steps, image=latent_image, generator=generator).images[0]
+        return seed, refined_image
+
+# Define the interface layout
+css = '''
+.gradio-container{max-width: 700px !important}
+h1{text-align:center}
+footer { visibility: hidden }
+'''
+
+examples = [
+    ["Image Generation", None, "A luxurious supercar with a unique design. The car should have a pearl white finish, and gold accents. 4k, realistic."],
+    ["Image Editing", "./supercar.png", "make it red"],
+    ["Image Editing", "./red_car.png", "add some snow"],
+    ["Image Generation", None, "An alien grasping a sign board containing the word 'ALIEN' with Neon Glow, neon, futuristic, neonpunk."],
+]
+
+# Define Gradio interface
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# Image Generation & Editing")
+
+    with gr.Row():
+        instruction = gr.Textbox(lines=1, label="Instruction")
+        generate_button = gr.Button("Run")
+
+    with gr.Row():
+        task = gr.Dropdown(["Image Generation", "Image Editing"], label="Task", value="Image Generation")
+        enhance_prompt = gr.Checkbox(label="Enhance prompt", value=False)
+        fast = gr.Checkbox(label="Fast Generation", value=True)
+    
+    with gr.Row():
+        input_image = gr.Image(label="Image", type='filepath')
+
+    with gr.Row():
+        guidance_scale = gr.Slider(6.0, 1.0, 15.0, label="Guidance Scale")
+        steps = gr.Slider(25, 1, 100, label="Steps")
+    
+    with gr.Accordion("Advanced options", open=False):
+        negative_prompt = gr.Text(value="deformed, distorted, blurry", label="Negative prompt")
+        width = gr.Slider(1024, 256, 2048, step=64, label="Width")
+        height = gr.Slider(1024, 256, 2048, step=64, label="Height")
+        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+        seed = gr.Number(value=2404, label="Seed")
+
+    gr.Examples(examples=examples, inputs=[task, input_image, instruction], fn=king, outputs=[seed, input_image])
+
+    generate_button.click(king, [task, input_image, instruction, negative_prompt, enhance_prompt, steps, randomize_seed, seed, width, height, guidance_scale, fast], [seed, input_image])
+
+demo.launch()