Upload app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline
 import logging
 from PIL import Image
 import base64
@@ -10,134 +10,64 @@ from io import BytesIO
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+logger.info("Starting model initialization...")
+
 try:
-    logger.info("Starting model initialization...")
-    # Initialize the FLUX model
-    pipe = DiffusionPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-schnell",
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/flux-1-schnell",
         torch_dtype=torch.float32,
         safety_checker=None,
         requires_safety_checker=False
     )
-    
-    # Force CPU usage and optimize memory
-    pipe = pipe.to("cpu")
-    pipe.enable_attention_slicing()
-    logger.info("Model initialization completed successfully")
+    pipe.enable_model_cpu_offload()
+    logger.info("Model initialized successfully")
 except Exception as e:
     logger.error(f"Error during model initialization: {str(e)}")
     raise
 
-def generate_image(prompt: str, negative_prompt: str, style: str) -> dict:
+def generate_image(prompt, negative_prompt="", style=""):
     try:
-        logger.info(f"Starting image generation with prompt: {prompt}, style: {style}")
-        
-        # Base negative prompt for better quality
-        base_negative = "ugly, deformed, noisy, blurry, distorted, grainy, text, watermark, signature, frame, deformed face"
-        negative_prompt = f"{base_negative}, {negative_prompt}" if negative_prompt else base_negative
-        
-        # Adjust the prompt based on style
-        if style == "Realistic":
-            prompt = f"ultra realistic, highly detailed photograph, professional photography, 8k uhd, muted colors, {prompt}"
-        elif style == "Artistic":
-            prompt = f"artistic, vibrant colors, stylized, digital art masterpiece, detailed, {prompt}"
-        elif style == "Abstract":
-            prompt = f"abstract, modern art, conceptual, artistic composition, cold color palette, {prompt}"
-        elif style == "Minimalist":
-            prompt = f"minimalist, clean lines, simple composition, elegant, muted colors, {prompt}"
+        # Enhance prompt based on style
+        if style:
+            style_prompts = {
+                "Realistic": "highly detailed, photorealistic, 8k resolution",
+                "Artistic": "artistic style, vibrant colors, creative composition",
+                "Abstract": "abstract art style, modern, conceptual",
+                "Minimalist": "minimalist style, clean lines, simple composition"
+            }
+            prompt = f"{prompt}, {style_prompts.get(style, '')}"
 
-        logger.info(f"Modified prompt: {prompt}")
-        logger.info(f"Negative prompt: {negative_prompt}")
-        
-        # Generate the image
-        with gr.Progress() as progress:
-            progress(0, desc="Starting generation...")
-            image = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=20,
-                width=512,
-                height=512,
-                guidance_scale=7.5
-            ).images[0]
-            progress(1, desc="Finalizing...")
-        
-        logger.info("Image generation completed")
-        
-        # Convert PIL image to base64
+        # Generate image
+        logger.info(f"Generating image with prompt: {prompt}")
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=20,
+            guidance_scale=7.5
+        ).images[0]
+
+        # Convert to base64
         buffered = BytesIO()
         image.save(buffered, format="PNG")
         img_str = base64.b64encode(buffered.getvalue()).decode()
         
-        return {"data": [img_str]}
-        
+        return {"image": f"data:image/png;base64,{img_str}"}
     except Exception as e:
-        error_msg = f"Error during image generation: {str(e)}"
-        logger.error(error_msg)
-        return {"error": error_msg}
+        logger.error(f"Error generating image: {str(e)}")
+        return {"error": str(e)}
 
 # Create Gradio interface
-with gr.Blocks(theme=gr.themes.Default()) as iface:
-    gr.Markdown(
-        """
-        # 🖼 Uniwalls AI Wallpaper Generator
-        Create beautiful wallpapers using FLUX AI. Simply describe what you want to see!
-        """
-    )
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Describe what you want to generate...",
-                lines=3
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                placeholder="What you don't want to see...",
-                lines=2
-            )
-            style = gr.Radio(
-                choices=["Realistic", "Artistic", "Abstract", "Minimalist"],
-                label="Style",
-                value="Realistic"
-            )
-            generate_btn = gr.Button("🎨 Generate Wallpaper")
-        
-        with gr.Column(scale=3):
-            output = gr.JSON(label="Generated Image")
-            
-    # Example prompts
-    gr.Examples(
-        examples=[
-            ["A stunning mountain landscape at golden hour, cinematic", "people, text, watermark", "Realistic"],
-            ["Ethereal waves of light and color in motion", "sharp edges, realistic", "Abstract"],
-            ["Zen garden with cherry blossoms, peaceful atmosphere", "cluttered, busy", "Minimalist"],
-            ["Magical crystal cave with glowing elements", "photorealistic, mundane", "Artistic"],
-        ],
-        inputs=[prompt, negative_prompt, style],
-        outputs=output,
-        fn=generate_image,
-        cache_examples=True,
-    )
-    
-    # Set up the generate button click event
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, negative_prompt, style],
-        outputs=output,
-    )
-    
-    gr.Markdown(
-        """
-        ### 💡 Tips for best results:
-        - Be specific in your descriptions
-        - Add details about lighting, atmosphere, and mood
-        - Use the negative prompt to remove unwanted elements
-        - Try different styles to find the perfect look
-        - Add keywords like "cinematic", "professional", or "masterpiece"
-        """
-    )
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Describe the wallpaper you want to generate..."),
+        gr.Textbox(label="Negative Prompt", placeholder="What you don't want to see in the image..."),
+        gr.Dropdown(choices=["Realistic", "Artistic", "Abstract", "Minimalist"], label="Style", value="Realistic")
+    ],
+    outputs=gr.Image(type="pil"),
+    title="FLUX AI Wallpaper Generator",
+    description="Generate unique wallpapers using FLUX AI. Enter a prompt and select a style to create your custom wallpaper."
+)
 
 # Launch with basic settings
 if __name__ == "__main__":