Create app.py

app.py

@@ -0,0 +1,210 @@
+import torch
+import datetime
+from typing import List, Dict, Optional
+from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
+from PIL import Image
+import gradio as gr
+from transformers import pipeline as hf_pipeline
+
+class StableDiffusionAgent:
+    def __init__(self, config: Optional[Dict] = None):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.default_config = {
+            "model": "stabilityai/stable-diffusion-2-1",
+            "safety_checker": True,
+            "max_resolution": 1024,
+            "art_styles": ["realistic", "anime", "cyberpunk", "watercolor", "pixel-art"],
+            "default_style": "realistic",
+            "memory_size": 10,
+            "prompt_enhancer": True
+        }
+        self.config = {**self.default_config, **(config or {})}
+        
+        self._initialize_models()
+        self.memory = []
+        self.user_profiles = {}
+        self.current_style = self.config["default_style"]
+
+    def _initialize_models(self):
+        """Load all required models"""
+        # Text-to-Image Pipeline
+        self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
+            self.config["model"],
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+            safety_checker=None if not self.config["safety_checker"] else None
+        ).to(self.device)
+        
+        self.sd_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+            self.sd_pipeline.scheduler.config
+        )
+        
+        if self.device == "cuda":
+            self.sd_pipeline.enable_xformers_memory_efficient_attention()
+            self.sd_pipeline.enable_attention_slicing()
+
+        # Prompt Enhancement Model
+        if self.config["prompt_enhancer"]:
+            self.prompt_pipeline = hf_pipeline(
+                "text2text-generation",
+                model="microsoft/Promptist"
+            )
+
+    def _enhance_prompt(self, prompt: str) -> str:
+        """Improve prompt using LLM"""
+        if self.config["prompt_enhancer"]:
+            try:
+                return self.prompt_pipeline(prompt, max_length=256)[0]["generated_text"]
+            except:
+                return prompt
+        return prompt
+
+    def _apply_style(self, prompt: str, style: str) -> str:
+        """Apply artistic style to prompt"""
+        style_templates = {
+            "anime": "anime style, vibrant colors, detailed line art",
+            "cyberpunk": "neon lights, cyberpunk style, rainy night, futuristic",
+            "watercolor": "watercolor painting, soft edges, artistic",
+            "pixel-art": "8-bit pixel art, retro gaming style"
+        }
+        return f"{prompt}, {style_templates.get(style, '')}"
+
+    def generate(
+        self,
+        user_id: str,
+        prompt: str,
+        negative_prompt: str = "",
+        style: Optional[str] = None,
+        **kwargs
+    ) -> Dict:
+        """Main generation method with user context"""
+        # Get user preferences
+        user_prefs = self.user_profiles.get(user_id, {})
+        
+        # Enhance prompt
+        enhanced_prompt = self._enhance_prompt(prompt)
+        
+        # Apply style
+        style = style or user_prefs.get("style", self.current_style)
+        final_prompt = self._apply_style(enhanced_prompt, style)
+        
+        # Generate image
+        results = self._generate_image(
+            prompt=final_prompt,
+            negative_prompt=negative_prompt,
+            **{**self._get_default_params(), **kwargs}
+        )
+        
+        # Update memory and user profile
+        self._update_memory(user_id, prompt, results)
+        return {
+            "images": results["images"],
+            "metadata": {
+                "enhanced_prompt": enhanced_prompt,
+                "style": style,
+                "seed": results["seed"],
+                "timestamp": datetime.datetime.now().isoformat()
+            }
+        }
+
+    def _generate_image(self, **kwargs) -> Dict:
+        """Low-level generation with safety checks"""
+        generator = torch.Generator(device=self.device)
+        seed = kwargs.pop("seed", None)
+        if seed is not None:
+            generator = generator.manual_seed(seed)
+            
+        results = self.sd_pipeline(**kwargs, generator=generator)
+        
+        # Filter NSFW content
+        safe_images = []
+        for i, img in enumerate(results.images):
+            if results.nsfw_content_detected and results.nsfw_content_detected[i]:
+                safe_images.append(self._create_black_image(kwargs["width"], kwargs["height"]))
+            else:
+                safe_images.append(img)
+                
+        return {
+            "images": safe_images,
+            "seed": seed or generator.initial_seed()
+        }
+
+    def _update_memory(self, user_id: str, prompt: str, results: Dict):
+        """Store generation history"""
+        self.memory.append({
+            "user_id": user_id,
+            "prompt": prompt,
+            "timestamp": datetime.datetime.now(),
+            "metadata": results["metadata"]
+        })
+        if len(self.memory) > self.config["memory_size"]:
+            self.memory.pop(0)
+
+    def _get_default_params(self):
+        return {
+            "height": 512,
+            "width": 512,
+            "num_images_per_prompt": 1,
+            "num_inference_steps": 50,
+            "guidance_scale": 7.5
+        }
+
+    def _create_black_image(self, width: int, height: int) -> Image.Image:
+        return Image.new("RGB", (width, height), (0, 0, 0))
+
+    # ----------- User Interaction Methods -----------
+    def set_style(self, user_id: str, style: str):
+        if style in self.config["art_styles"]:
+            self.user_profiles.setdefault(user_id, {})["style"] = style
+            return f"Style set to {style}"
+        return f"Invalid style. Available styles: {', '.join(self.config['art_styles'])}"
+
+    def get_history(self, user_id: str) -> List[Dict]:
+        return [entry for entry in self.memory if entry["user_id"] == user_id]
+
+# ------------------ Gradio Interface ------------------
+def create_web_interface(agent: StableDiffusionAgent):
+    css = """
+    .gradio-container {max-width: 900px!important}
+    .output-image img {box-shadow: 0 4px 8px rgba(0,0,0,0.1)}
+    """
+    
+    with gr.Blocks(css=css) as interface:
+        gr.Markdown("# 🎨 AI Art Generator Agent")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                user_id = gr.Textbox(label="User ID", placeholder="Enter unique identifier")
+                prompt = gr.Textbox(label="Prompt", lines=3)
+                negative_prompt = gr.Textbox(label="Negative Prompt")
+                style = gr.Dropdown(agent.config["art_styles"], label="Art Style")
+                generate_btn = gr.Button("Generate", variant="primary")
+                
+            with gr.Column(scale=1):
+                output_image = gr.Image(label="Generated Art", elem_classes=["output-image"])
+                meta_info = gr.JSON(label="Generation Metadata")
+        
+        with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row():
+                steps = gr.Slider(10, 100, value=50, label="Steps")
+                guidance = gr.Slider(1.0, 20.0, value=7.5, label="Guidance Scale")
+                seed = gr.Number(label="Seed (optional)")
+        
+        generate_btn.click(
+            fn=lambda *args: agent.generate(*args),
+            inputs=[user_id, prompt, negative_prompt, style, steps, guidance, seed],
+            outputs=[output_image, meta_info]
+        )
+        
+    return interface
+
+if __name__ == "__main__":
+    # Initialize agent
+    config = {
+        "prompt_enhancer": True,
+        "art_styles": ["realistic", "anime", "cyberpunk", "watercolor"]
+    }
+    agent = StableDiffusionAgent(config)
+    
+    # Launch Gradio interface
+    interface = create_web_interface(agent)
+    interface.launch(server_port=7860, share=True)