Create app.py

app.py

@@ -0,0 +1,194 @@
+import spaces
+import os
+import torch
+import random
+from huggingface_hub import snapshot_download
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL
+from diffusers import EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DPMSolverSDEScheduler
+import gradio as gr
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+# Download the model files
+ckpt_dir = snapshot_download(repo_id="John6666/pony-realism-v21main-sdxl")
+
+# Load the models
+vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.float16)
+
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    ckpt_dir,
+    vae=vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+)
+pipe = pipe.to("cuda")
+
+# Define samplers
+samplers = {
+    "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config),
+    "DPM++ 2M": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, algorithm_type="dpmsolver++", use_karras_sigmas=True),
+    "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
+}
+
+DEFAULT_POSITIVE_PREFIX = "score_9, score_8_up, score_7_up, BREAK,"
+DEFAULT_POSITIVE_SUFFIX = "(masterpiece), best quality, very aesthetic, perfect face"
+DEFAULT_NEGATIVE_PREFIX = "score_1, score_2, score_3, text"
+DEFAULT_NEGATIVE_SUFFIX = "nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"
+
+# Initialize Florence model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
+florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
+
+# Prompt Enhancer
+enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
+enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
+
+# Florence caption function
+def florence_caption(image):
+    # Convert image to PIL if it's not already
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    inputs = florence_processor(text="<DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
+    generated_ids = florence_model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        early_stopping=False,
+        do_sample=False,
+        num_beams=3,
+    )
+    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = florence_processor.post_process_generation(
+        generated_text,
+        task="<DETAILED_CAPTION>",
+        image_size=(image.width, image.height)
+    )
+    return parsed_answer["<DETAILED_CAPTION>"]
+
+# Prompt Enhancer function
+def enhance_prompt(input_prompt, model_choice):
+    if model_choice == "Medium":
+        result = enhancer_medium("Enhance the description: " + input_prompt)
+        enhanced_text = result[0]['summary_text']
+    else:  # Long
+        result = enhancer_long("Enhance the description: " + input_prompt)
+        enhanced_text = result[0]['summary_text']
+    
+    return enhanced_text
+
+@spaces.GPU(duration=120)
+def generate_image(additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps, guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer, input_image=None, progress=gr.Progress(track_tqdm=True)):
+    if use_random_seed:
+        seed = random.randint(0, 2**32 - 1)
+    else:
+        seed = int(seed)  # Ensure seed is an integer
+    
+    # Set the scheduler based on the selected sampler
+    pipe.scheduler = samplers[sampler]
+    
+    # Set clip skip
+    pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)
+    
+    # Start with the default positive prompt prefix
+    full_positive_prompt = DEFAULT_POSITIVE_PREFIX
+
+    # Add Florence-2 caption if enabled and image is provided
+    if use_florence2 and input_image is not None:
+        florence2_caption = florence_caption(input_image)
+        florence2_caption = florence2_caption.lower().replace('.', ',')
+        additional_positive_prompt = f"{florence2_caption}, {additional_positive_prompt}" if additional_positive_prompt else florence2_caption
+
+    # Enhance only the additional positive prompt if enhancers are enabled
+    if additional_positive_prompt:
+        enhanced_prompt = additional_positive_prompt
+        if use_medium_enhancer:
+            medium_enhanced = enhance_prompt(enhanced_prompt, "Medium")
+            medium_enhanced = medium_enhanced.lower().replace('.', ',')
+            enhanced_prompt = f"{enhanced_prompt}, {medium_enhanced}"
+        if use_long_enhancer:
+            long_enhanced = enhance_prompt(enhanced_prompt, "Long")
+            long_enhanced = long_enhanced.lower().replace('.', ',')
+            enhanced_prompt = f"{enhanced_prompt}, {long_enhanced}"
+        full_positive_prompt += f"{enhanced_prompt}"
+
+    # Add the default positive suffix
+    full_positive_prompt += f", {DEFAULT_POSITIVE_SUFFIX}"
+    
+    # Combine default negative prompt with additional negative prompt
+    full_negative_prompt = f"{DEFAULT_NEGATIVE_PREFIX}, {additional_negative_prompt}, {DEFAULT_NEGATIVE_SUFFIX}" if additional_negative_prompt else f"{DEFAULT_NEGATIVE_PREFIX}, {DEFAULT_NEGATIVE_SUFFIX}"
+    
+    try:
+        image = pipe(
+            prompt=full_positive_prompt,
+            negative_prompt=full_negative_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            generator=torch.Generator(pipe.device).manual_seed(seed)
+        ).images
+        return image, seed, full_positive_prompt
+    except Exception as e:
+        print(f"Error during image generation: {str(e)}")
+        return None, seed, full_positive_prompt
+
+# Gradio interface
+with gr.Blocks(theme='bethecloud/storj_theme') as demo:
+    gr.HTML("""
+    <h1 align="center">Pony Realism v21 SDXL - Text-to-Image Generation</h1>
+    <p align="center">
+    <a href="https://huggingface.co/John6666/pony-realism-v21main-sdxl/" target="_blank">[HF Model Page]</a>
+    <a href="https://civitai.com/models/372465/pony-realism" target="_blank">[civitai Model Page]</a>
+    <a href="https://huggingface.co/microsoft/Florence-2-base" target="_blank">[Florence-2 Model]</a>
+    <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance-Long" target="_blank">[Prompt Enhancer Long]</a>
+    <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance" target="_blank">[Prompt Enhancer Medium]</a>
+    </p>
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            positive_prompt = gr.Textbox(label="Positive Prompt", placeholder="Add your positive prompt here")
+            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Add your negative prompt here")
+            
+            with gr.Accordion("Advanced settings", open=False):
+                height = gr.Slider(512, 2048, 1024, step=64, label="Height")
+                width = gr.Slider(512, 2048, 1024, step=64, label="Width")
+                num_inference_steps = gr.Slider(20, 50, 30, step=1, label="Number of Inference Steps")
+                guidance_scale = gr.Slider(1, 20, 6, step=0.1, label="Guidance Scale")
+                num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt")
+                use_random_seed = gr.Checkbox(label="Use Random Seed", value=True)
+                seed = gr.Number(label="Seed", value=0, precision=0)
+                sampler = gr.Dropdown(label="Sampler", choices=list(samplers.keys()), value="DPM++ SDE Karras")
+                clip_skip = gr.Slider(1, 4, 2, step=1, label="Clip skip")
+            
+            with gr.Accordion("Captioner and Enhancers", open=False):
+                input_image = gr.Image(label="Input Image for Florence-2 Captioner")
+                use_florence2 = gr.Checkbox(label="Use Florence-2 Captioner", value=False)
+                use_medium_enhancer = gr.Checkbox(label="Use Medium Prompt Enhancer", value=False)
+                use_long_enhancer = gr.Checkbox(label="Use Long Prompt Enhancer", value=False)
+            
+            generate_btn = gr.Button("Generate Image")
+
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(label="Result", elem_id="gallery", show_label=False)
+            seed_used = gr.Number(label="Seed Used")
+            full_prompt_used = gr.Textbox(label="Full Positive Prompt Used")
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[
+            positive_prompt, negative_prompt, height, width, num_inference_steps,
+            guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
+            clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer, input_image
+        ],
+        outputs=[output_gallery, seed_used, full_prompt_used]
+    )
+
+demo.launch(debug=True)