Create app.py

app.py

@@ -0,0 +1,200 @@
+import spaces
+import gradio as gr
+import torch
+from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, pipeline
+from diffusers import StableDiffusion3Pipeline
+import re
+import random
+import numpy as np
+import os
+from huggingface_hub import snapshot_download
+
+# Initialize models
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16
+
+huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+
+model_path = snapshot_download(
+    repo_id="stabilityai/stable-diffusion-3-medium", 
+    repo_type="model", 
+    ignore_patterns=["*.md", "*..gitattributes"],
+    local_dir="SD3",
+    token=huggingface_token, # type a new token-id.
+    )
+
+# VLM Captioner
+vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to(device).eval()
+vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner")
+
+# 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)
+
+# SD3
+sd3_pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1344
+
+# VLM Captioner function
+def create_captions_rich(image):
+    prompt = "caption en"
+    model_inputs = vlm_processor(text=prompt, images=image, return_tensors="pt").to(device)
+    input_len = model_inputs["input_ids"].shape[-1]
+
+    with torch.inference_mode():
+        generation = vlm_model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
+        generation = generation[0][input_len:]
+        decoded = vlm_processor.decode(generation, skip_special_tokens=True)
+
+    return modify_caption(decoded)
+
+# Helper function for caption modification
+def modify_caption(caption: str) -> str:
+    prefix_substrings = [
+        ('captured from ', ''),
+        ('captured at ', '')
+    ]
+    pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
+    replacers = {opening: replacer for opening, replacer in prefix_substrings}
+    
+    def replace_fn(match):
+        return replacers[match.group(0)]
+    
+    return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
+
+# 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']
+        
+        pattern = r'^.*?of\s+(.*?(?:\.|$))'
+        match = re.match(pattern, enhanced_text, re.IGNORECASE | re.DOTALL)
+        
+        if match:
+            remaining_text = enhanced_text[match.end():].strip()
+            modified_sentence = match.group(1).capitalize()
+            enhanced_text = modified_sentence + ' ' + remaining_text
+    else:  # Long
+        result = enhancer_long("Enhance the description: " + input_prompt)
+        enhanced_text = result[0]['summary_text']
+    
+    return enhanced_text
+
+# SD3 Generation function
+def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    generator = torch.Generator().manual_seed(seed)
+    
+    image = sd3_pipe(
+        prompt=prompt, 
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale, 
+        num_inference_steps=num_inference_steps, 
+        width=width, 
+        height=height,
+        generator=generator
+    ).images[0]
+    
+    return image, seed
+
+# Gradio Interface
+@spaces.GPU
+def process_workflow(image, text_prompt, use_vlm, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    if use_vlm and image is not None:
+        prompt = create_captions_rich(image)
+    else:
+        prompt = text_prompt
+    
+    if use_enhancer:
+        prompt = enhance_prompt(prompt, model_choice)
+    
+    generated_image, used_seed = generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps)
+    
+    return generated_image, prompt, used_seed
+
+
+css = """
+body {
+    font-family: 'Arial', sans-serif;
+    background-color: #f0f4f8;
+}
+.container {
+    max-width: 800px;
+    margin: 0 auto;
+    padding: 20px;
+    background-color: white;
+    border-radius: 10px;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+}
+h1 {
+    color: #2c3e50;
+    text-align: center;
+    margin-bottom: 20px;
+}
+.input-box, .output-box {
+    border: 1px solid #bdc3c7;
+    border-radius: 5px;
+    padding: 10px;
+}
+.input-box:focus, .output-box:focus {
+    border-color: #3498db;
+    box-shadow: 0 0 5px rgba(52, 152, 219, 0.5);
+}
+.submit-btn {
+    background-color: #2980b9;
+    color: white;
+    border: none;
+    padding: 10px 20px;
+    border-radius: 5px;
+    cursor: pointer;
+    transition: background-color 0.3s;
+}
+.submit-btn:hover {
+    background-color: #3498db;
+}
+"""
+
+    
+# Gradio Interface
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# SD3 Image Generator + VLM Captioner + Prompt Enhancer")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image for VLM")
+            text_prompt = gr.Textbox(label="Text Prompt")
+            use_vlm = gr.Checkbox(label="Use VLM Captioner", value=False)
+            use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False)
+            model_choice = gr.Radio(["Medium", "Long"], label="Enhancer Model", value="Long")
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                negative_prompt = gr.Textbox(label="Negative Prompt")
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+                guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=5.0)
+                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=28)
+            
+            generate_btn = gr.Button("Generate Image")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+            final_prompt = gr.Textbox(label="Final Prompt Used")
+            used_seed = gr.Number(label="Seed Used")
+    
+    generate_btn.click(
+        fn=process_workflow,
+        inputs=[
+            input_image, text_prompt, use_vlm, use_enhancer, model_choice,
+            negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps
+        ],
+        outputs=[output_image, final_prompt, used_seed]
+    )
+
+demo.launch()