Create image_generation.py

image_generation.py

@@ -0,0 +1,92 @@
+import replicate
+from PIL import Image
+import io
+import requests
+import base64
+
+def generate_image(
+    prompt,
+    num_steps=30,
+    guidance_scale=7.5,
+    aspect_ratio="1:1",
+    replicate_api_key=None,
+    lora_url=None,
+    negative_prompt=None
+):
+    """
+    Generate an image using Stable Diffusion via Replicate API
+    
+    Args:
+        prompt (str): The text prompt for image generation
+        num_steps (int): Number of inference steps
+        guidance_scale (float): Guidance scale for generation
+        aspect_ratio (str): Desired aspect ratio ("1:1", "16:9", "3:2", etc.)
+        replicate_api_key (str): API key for Replicate
+        lora_url (str, optional): URL to LoRA weights
+        negative_prompt (str, optional): Negative prompt for generation
+    """
+    try:
+        if not replicate_api_key:
+            return None, "Please provide a Replicate API key"
+            
+        # Set up aspect ratio dimensions
+        aspect_ratios = {
+            "1:1": (512, 512),
+            "16:9": (912, 512),
+            "3:2": (768, 512),
+            "2:3": (512, 768),
+            "4:5": (512, 640),
+            "5:4": (640, 512)
+        }
+        width, height = aspect_ratios.get(aspect_ratio, (512, 512))
+        
+        # Configure model parameters
+        model_params = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt or "ugly, blurry, low quality, distorted, deformed",
+            "num_inference_steps": num_steps,
+            "guidance_scale": guidance_scale,
+            "width": width,
+            "height": height,
+            "scheduler": "DPMSolverMultistep",  # You can experiment with different schedulers
+            "num_outputs": 1,
+        }
+        
+        # Add LoRA if specified
+        if lora_url:
+            model_params["lora_urls"] = lora_url
+        
+        # Set API key
+        client = replicate.Client(api_token=replicate_api_key)
+        
+        # Run the model
+        # Using SDXL model for better quality
+        output = client.run(
+            "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b",
+            input=model_params
+        )
+        
+        # Get the image URL from output
+        if output and len(output) > 0:
+            image_url = output[0]
+            
+            # Download and convert to PIL Image
+            response = requests.get(image_url)
+            if response.status_code == 200:
+                image = Image.open(io.BytesIO(response.content))
+                return image, "Success"
+            else:
+                return None, f"Failed to download image: {response.status_code}"
+        else:
+            return None, "No image generated"
+        
+    except Exception as e:
+        return None, f"Error generating image: {str(e)}"
+
+def encode_image_to_base64(image):
+    """Helper function to convert PIL Image to base64 string"""
+    if isinstance(image, Image.Image):
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+    return None