Update app.py

app.py

@@ -1,76 +1,120 @@
-import os
+import gradio as gr
 import torch
+import numpy as np
 from diffusers import DiffusionPipeline
-import gradio as gr
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from functools import lru_cache
+from PIL import Image
+from torchvision import transforms
+
 
-# Load the base model and apply the LoRA weights for super realism
+# Cache pipeline loading to improve performance
+@lru_cache(maxsize=1)
 def load_pipeline():
+    # Load base model
     base_model = "black-forest-labs/FLUX.1-dev"
-    lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
-    trigger_word = "Super Realism"  # Recommended trigger word
-
-    
-
     pipe = DiffusionPipeline.from_pretrained(
         base_model,
-        torch_dtype=torch.bfloat16,
-       
+        torch_dtype=torch.bfloat16
     )
-
-    # Load the LoRA weights into the pipeline
+    
+    # Load LoRA weights
+    lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
     pipe.load_lora_weights(lora_repo)
+    
+    # Load safety checker
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
+    
+    # Optimizations
+    pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to("cuda")
+    
+    return pipe, safety_checker, feature_extractor
 
-    # Use GPU if available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe.to(device)
-    return pipe
-
-# Instantiate the pipeline once on Space startup
-pipe = load_pipeline()
-
-# Define a function for image generation
-def generate_image(prompt, seed, width, height, guidance_scale, randomize_seed):
-    # If randomize_seed is selected, allow the model to generate a random seed
-    if randomize_seed:
-        seed = None
+pipe, safety_checker, feature_extractor = load_pipeline()
 
-    # Ensure the prompt includes realism trigger words if needed
-    if "realistic" not in prompt.lower() and "realism" not in prompt.lower():
-        prompt += " realistic, realism"
+def generate_image(
+    prompt,
+    seed=42,
+    width=1024,
+    height=1024,
+    guidance_scale=6,
+    steps=28,
+    progress=gr.Progress()
+):
+    try:
+        progress(0, desc="Initializing...")
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        
+        # Auto-add trigger words
+        if "super realism" not in prompt.lower():
+            prompt = f"Super Realism, {prompt}"
+        
+        # Create callback for progress updates
+        def update_progress(step, _, __):
+            progress((step + 1) / steps, desc="Generating image...")
+        
+        # Generate image
+        with torch.inference_mode():
+            image = pipe(
+                prompt=prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=steps,
+                generator=generator,
+                callback=update_progress
+            ).images[0]
+        
+        # Safety check
+        progress(1, desc="Safety checking...")
+        safety_input = feature_extractor(image, return_tensors="pt")
+        np_image = np.array(image)
+        safety_result = safety_checker(
+            images=[np_image], 
+            clip_input=safety_input.pixel_values
+        )
+        
+        if safety_result.nsfw[0]:
+            return Image.new("RGB", (512, 512)), "NSFW content detected"
+        
+        return image, "Generation successful"
+    
+    except Exception as e:
+        return Image.new("RGB", (512, 512)), f"Error: {str(e)}"
 
-    # Generate the image
-    output = pipe(
-        prompt=prompt,
-        seed=seed,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale
+# Create Gradio interface with rate limiting
+with gr.Blocks() as app:
+    gr.Markdown("# Flux Super Realism Generator")
+    
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", value="A portrait of a person")
+            seed = gr.Slider(0, 1000, value=42, label="Seed")
+            width = gr.Slider(512, 2048, value=1024, label="Width")
+            height = gr.Slider(512, 2048, value=1024, label="Height")
+            guidance = gr.Slider(1, 20, value=6, label="Guidance Scale")
+            steps = gr.Slider(10, 100, value=28, label="Steps")
+            submit = gr.Button("Generate")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Result", type="pil")
+            status = gr.Textbox(label="Status")
+    
+    submit.click(
+        generate_image,
+        inputs=[prompt, seed, width, height, guidance, steps],
+        outputs=[output_image, status]
     )
-    return output.images[0]
-
-# Set up Gradio interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(
-            lines=2, 
-            label="Prompt", 
-            placeholder="Enter your prompt, e.g., 'A tiny astronaut hatching from an egg on the moon, 4k, planet theme'"
-        ),
-        gr.Slider(0, 10000, step=1, value=0, label="Seed (0 for random)"),
-        gr.Slider(256, 1024, step=64, value=1024, label="Width"),
-        gr.Slider(256, 1024, step=64, value=1024, label="Height"),
-        gr.Slider(1, 20, step=0.5, value=6, label="Guidance Scale"),
-        gr.Checkbox(value=True, label="Randomize Seed")
-    ],
-    outputs=gr.Image(type="pil", label="Generated Image"),
-    title="Flux Super Realism LoRA Demo",
-    description=(
-        "This demo uses the Flux Super Realism LoRA model for ultra-realistic image generation. "
-        "You can use the trigger word 'Super Realism' (recommended) along with other realism-related words "
-        "to guide the generation process."
-    ),
-)
+    
+    # Rate limiting example (1 request every 30 seconds)
+    app.queue(concurrency_count=1, max_size=3).launch()
 
-if __name__ == "__main__":
-    iface.launch(share=False) 
+# For multiple GPU support (advanced)
+# pipe.enable_model_cpu_offload()
+# pipe.enable_sequential_cpu_offload()