00Boobs00's picture
🎨 Redesign from AnyCoder (#1)
b391a1c verified
"""
Omni Image Editor - Redesigned
A modern, mobile-first Gradio 6 interface for image editing and generation.
"""
import gradio as gr
from pipeline import OmniMMDitV2Pipeline
from nfsw_filter import NSFWDetector
# Initialize Pipeline and Safety Filter
# Note: Ensure you have the required models downloaded or accessible
try:
pipeline = OmniMMDitV2Pipeline(
model=None, # Placeholder: Load actual model weights here
vae=None, # Placeholder: Load actual VAE here
text_encoder=None, # Placeholder
tokenizer=None, # Placeholder
scheduler=None, # Placeholder
)
safety_checker = NSFWDetector()
except Exception as e:
print(f"Warning: Models could not be initialized: {e}")
# Create dummy instances for UI demonstration if models fail to load
pipeline = None
safety_checker = None
def check_safety(image):
"""Run NSFW check on the image."""
if safety_checker is None:
return image, "Safety check unavailable."
try:
# Assuming PIL image input
label, _ = safety_checker.predict_from_pil(image)
if label == "nsfw":
return None, "⚠️ NSFW content detected. Image blocked."
return image, "✅ Image is safe."
except Exception as e:
return image, f"Error checking safety: {e}"
def edit_image(input_image, prompt, negative_prompt, guidance_scale, num_steps):
"""Process image editing request."""
if input_image is None:
return None, "Please upload an image first."
if pipeline is None:
return None, "Model not loaded. Check server logs."
# 1. Safety Check
safe_image, safety_msg = check_safety(input_image)
if safe_image is None:
return None, safety_msg
# 2. Run Inference
try:
# Note: This is a mockup of the pipeline call.
# Adjust parameters based on actual OmniMMDitV2Pipeline implementation
result = pipeline(
prompt=prompt,
input_images=[input_image], # Image-to-Image
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=int(num_steps),
height=1024,
width=1024
)
# Extract image from result (assuming pipeline returns dict or object)
if hasattr(result, 'images'):
output_image = result.images[0]
else:
output_image = result[0]
return output_image, f"Success! {safety_msg}"
except Exception as e:
return None, f"Error during generation: {str(e)}"
# --- Gradio 6 Interface ---
# Custom Theme: Modern, Soft, Mobile-Friendly
custom_theme = gr.themes.Soft(
primary_hue="indigo",
secondary_hue="cyan",
neutral_hue="slate",
font=gr.themes.GoogleFont("Inter"),
text_size="lg",
spacing_size="lg",
radius_size="lg"
).set(
# Button styling
button_primary_background_fill="*primary_500",
button_primary_background_fill_hover="*primary_600",
button_primary_text_color="white",
button_secondary_background_fill="*neutral_100",
button_secondary_background_fill_hover="*neutral_200",
# Container styling
block_background_fill="*background_fill_primary",
block_border_width="0px",
block_shadow="*shadow_drop_lg",
# Typography
block_title_text_weight="600",
block_title_text_size="*text_xl",
)
with gr.Blocks() as demo:
# Header / Navbar
gr.HTML("""
<div style="display: flex; align-items: center; justify-content: space-between; padding: 1rem 0;">
<div>
<h1 style="margin: 0; font-size: 1.5rem; font-weight: 700; color: #4F46E5;">Omni Editor</h1>
<p style="margin: 0; color: #64748B;">Next-Gen Image Editing</p>
</div>
<a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="text-decoration: none; color: #4F46E5; font-weight: 500; font-size: 0.9rem;">
Built with anycoder ↗
</a>
</div>
""")
with gr.Row(equal_height=True):
# --- Sidebar (Controls) ---
with gr.Column(scale=1, min_width=300):
gr.Markdown("### 🎨 Controls")
with gr.Accordion("⚙️ Generation Settings", open=True):
prompt = gr.Textbox(
label="Prompt",
placeholder="Describe the changes you want...",
lines=3,
show_label=False
)
negative_prompt = gr.Textbox(
label="Negative Prompt",
placeholder="Things to avoid...",
lines=2,
show_label=False
)
with gr.Row():
guidance = gr.Slider(
minimum=1.0,
maximum=20.0,
value=7.5,
step=0.5,
label="Guidance Scale"
)
steps = gr.Slider(
minimum=10,
maximum=50,
value=25,
step=1,
label="Steps"
)
generate_btn = gr.Button("✨ Generate Edit", variant="primary", size="lg")
status_msg = gr.Textbox(label="Status", interactive=False, visible=False)
gr.Markdown("""
---
### 💡 Tips
- Upload a clear image for best results.
- Be descriptive in your prompt.
- Use negative prompts to remove artifacts.
""")
# --- Main Content (Canvas) ---
with gr.Column(scale=2):
gr.Markdown("### 🖼️ Workspace")
with gr.Row():
input_image = gr.Image(
label="Input",
type="pil",
sources=["upload", "clipboard", "webcam"],
interactive=True
)
output_image = gr.Image(
label="Result",
type="pil",
interactive=False
)
# --- Event Listeners ---
generate_btn.click(
fn=edit_image,
inputs=[input_image, prompt, negative_prompt, guidance, steps],
outputs=[output_image, status_msg],
api_visibility="public"
).then(
# Hide status after success if desired, or keep it
lambda: gr.Textbox(visible=True),
outputs=[status_msg]
)
# --- Launch Configuration ---
demo.launch(
theme=custom_theme,
footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
)