Update app.py

app.py

@@ -1,27 +1,194 @@
 import gradio as gr
-from diffusers import AutoPipelineForText2Image
+import spaces
 import torch
+import os
+from compel import Compel, ReturnedEmbeddingsType
+from diffusers import DiffusionPipeline
 
-# Memuat model text-to-image
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipeline = AutoPipelineForText2Image.from_pretrained(
-    "stable-diffusion-v1-5/stable-diffusion-v1-5", 
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-).to(device)
-
-def generate_image(prompt):
-    # Menghasilkan gambar dari teks
-    image = pipeline(prompt).images[0]
-    return image
-
-# Membuat Antarmuka Gradio
-demo = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(label="Input your prompt"),
-    outputs=gr.Image(label="Image Result"),
-    title="AI Text-to-Image Generator",
-    description="Input Text via Stable Diffusion."
+# Load model
+model_name = os.environ.get('MODEL_NAME', 'UnfilteredAI/NSFW-gen-v2.1')
+pipe = DiffusionPipeline.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16
 )
+pipe.to('cuda')
 
-if __name__ == "__main__":
-    demo.launch()
+# Compel setup
+compel = Compel(
+  tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+  text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+  requires_pooled=[False, True]
+)
+
+# Default negative prompt
+default_negative_prompt = "(low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn, (deformed | distorted | disfigured:1.3), bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers:1.4, disconnected limbs, blurry, amputation."
+
+# Example prompts
+example_prompts = [
+    ["a beautiful 19 aged rusian girl in a summer dress at the beach, golden sunset, professional photography, 8k", default_negative_prompt, 40, 7.5, 1024, 1024, 4],
+    ["a beautiful japan woman in a vikini at the beach, noon sunset, professional photography, 8k", default_negative_prompt, 45, 7.0, 1024, 1024, 4],
+]
+
+# Image generation function
+@spaces.GPU(duration=120)
+def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples, progress=gr.Progress()):
+    progress(0, desc="Preparing")
+    embeds, pooled = compel(prompt)
+    neg_embeds, neg_pooled = compel(negative_prompt)
+    
+    progress(0.1, desc="Generating images")
+    
+    # Define proper callback for step end
+    def callback_on_step_end(pipe, i, t, callback_kwargs):
+        progress((i + 1) / num_inference_steps)
+        return callback_kwargs
+    
+    images = pipe(
+        prompt_embeds=embeds,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=neg_embeds,
+        negative_pooled_prompt_embeds=neg_pooled,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        width=width,
+        height=height,
+        num_images_per_prompt=num_samples,
+        callback_on_step_end=callback_on_step_end
+    ).images
+    
+    return images
+
+# CSS styles
+css = """
+.gallery-item {
+    transition: transform 0.2s;
+    box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+    border-radius: 10px;
+}
+.gallery-item:hover {
+    transform: scale(1.03);
+    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
+}
+.container {
+    max-width: 1200px;
+    margin: auto;
+}
+.header {
+    text-align: center;
+    margin-bottom: 2rem;
+    padding: 1rem;
+    background: linear-gradient(90deg, rgba(76,0,161,0.8) 0%, rgba(28,110,164,0.8) 100%);
+    border-radius: 10px;
+    color: white;
+}
+.slider-container {
+    background-color: #f5f5f5;
+    padding: 1rem;
+    border-radius: 10px;
+    margin-bottom: 1rem;
+}
+.prompt-container {
+    background-color: #f0f8ff;
+    padding: 1rem;
+    border-radius: 10px;
+    margin-bottom: 1rem;
+    border: 1px solid #d0e8ff;
+}
+.examples-header {
+    background: linear-gradient(90deg, rgba(41,128,185,0.7) 0%, rgba(142,68,173,0.7) 100%);
+    color: white;
+    padding: 0.5rem;
+    border-radius: 8px;
+    text-align: center;
+    margin-bottom: 0.5rem;
+}
+"""
+
+# Gradio interface
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    gr.HTML("""
+        <style>
+            .gradio-container {
+                background: linear-gradient(135deg, #fef9f3 0%, #f0e6fa 50%, #e6f0fa 100%) !important;
+            }
+            footer {display: none !important;}
+        </style>
+
+        <div style="text-align: center; margin-bottom: 20px;">
+            <h1 style="color: #6b5b7a; font-size: 2.2rem; font-weight: 700; margin-bottom: 0.3rem;">
+                🎨 Unfiltered AI NSFW Image Generator
+            </h1>
+
+            <p style="color: #8b7b9b; font-size: 1rem;">
+                Enter creative prompts and generate high-quality images.
+            </p>
+
+            <div style="margin-top: 12px; display: flex; justify-content: center; gap: 12px;">
+                <a href="https://huggingface.co/spaces/Heartsync/FREE-NSFW-HUB" target="_blank">
+                    <img src="https://img.shields.io/static/v1?label=FREE&message=NSFW%20HUB&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
+                </a>
+                <a href="https://www.humangen.ai" target="_blank">
+                  <img src="https://img.shields.io/static/v1?label=100% FREE&message=AI%20Playground&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="badge">
+                </a>                  
+                <a href="https://ginigen.ai/en" target="_blank">
+                    <img src="https://img.shields.io/static/v1?label=Powered%20by&message=Hogwarts%20BANANA&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
+                </a>
+            </div>
+        </div>
+    """)
+
+
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Group(elem_classes="prompt-container"):
+                prompt = gr.Textbox(label="Prompt", placeholder="Describe your desired image...", lines=3)
+                negative_prompt = gr.Textbox(
+                    label="Negative Prompt", 
+                    value=default_negative_prompt,
+                    lines=3
+                )
+            
+            with gr.Group(elem_classes="slider-container"):
+                with gr.Row():
+                    with gr.Column():
+                        steps = gr.Slider(minimum=20, maximum=100, value=60, step=1, label="Inference Steps (Quality)", info="Higher values improve quality (longer generation time)")
+                        guidance = gr.Slider(minimum=1, maximum=15, value=7, step=0.1, label="Guidance Scale (Creativity)", info="Lower values create more creative results")
+                    
+                    with gr.Column():
+                        with gr.Row():
+                            width = gr.Slider(minimum=512, maximum=1536, value=1024, step=128, label="Width")
+                            height = gr.Slider(minimum=512, maximum=1536, value=1024, step=128, label="Height")
+                        
+                        num_samples = gr.Slider(minimum=1, maximum=8, value=4, step=1, label="Number of Images", info="Number of images to generate at once")
+            
+            generate_btn = gr.Button("🚀 Generate Images", variant="primary", size="lg")
+        
+        with gr.Column(scale=3):
+            output_gallery = gr.Gallery(label="Generated Images", elem_classes="gallery-item", columns=2, object_fit="contain", height=650)
+    
+    gr.HTML("""<div class="examples-header"><h3>✨ Example Prompts</h3></div>""")
+    gr.Examples(
+        examples=example_prompts,
+        inputs=[prompt, negative_prompt, steps, guidance, width, height, num_samples],
+        outputs=output_gallery,
+        fn=generate,
+        cache_examples=True,
+    )
+    
+    # Event connections
+    generate_btn.click(
+        fn=generate,
+        inputs=[prompt, negative_prompt, steps, guidance, width, height, num_samples],
+        outputs=output_gallery
+    )
+    
+    gr.HTML("""
+    <div style="text-align: center; margin-top: 20px; padding: 10px; background-color: #f0f0f0; border-radius: 10px;">
+        <p>💡 Tip: For high-quality images, use detailed prompts and higher inference steps.</p>
+        <p>Example: Add quality terms like "professional photography, 8k, highly detailed, sharp focus, HDR" to your prompts.</p>
+    </div>
+    """)
+
+demo.launch()