""" title: Uncensored FLUX.1 Image Gen author: antonymajar version: 0.1 license: MIT description: Generate uncensored images using flux.1 with novita.ai # Instructions Import the tool Get an api key from novita.ai they will give you a free credit voucher ( use my affiliate link if you'd like to support me :) https://novita.ai/?ref=oty4ndc&utm_source=affiliate ) Set the api key by clicking the tool settings in your workspace """ import base64 import os import random from typing import Any, Dict, Generator, Iterator, List, Union import requests from open_webui.utils.misc import get_last_user_message from pydantic import BaseModel, Field class Pipe: """ Class representing the FLUX.1 Schnell image generation provider. """ class Valves(BaseModel): """ Pydantic model for storing API configuration. """ NOVITA_API_KEY: str = Field( default="", description="Your API Key for Novita.ai" ) NOVITA_API_BASE_URL: str = Field( default="https://api.novita.ai", description="Base URL for the Novita API" ) def __init__(self): """ Initialize the Pipe class with default values and environment variables. """ self.type = "manifold" self.id = "FLUX_SCHNELL" self.name = "FLUX.1 Schnell: " self.valves = self.Valves( NOVITA_API_KEY=os.getenv("NOVITA_API_KEY", ""), NOVITA_API_BASE_URL=os.getenv( "NOVITA_API_BASE_URL", "https://api.novita.ai" ), ) def get_image_from_url(self, url: str) -> str: """ Download and convert an image URL to base64 format. Args: url (str): The image URL Returns: str: Base64-encoded image data with content type """ response = requests.get(url) response.raise_for_status() content_type = response.headers.get("Content-Type", "image/jpeg") image_data = base64.b64encode(response.content).decode("utf-8") return f"data:{content_type};base64,{image_data}" def non_stream_response( self, headers: Dict[str, str], payload: Dict[str, Any] ) -> str: """ Process a non-streaming image generation request. Args: headers (Dict[str, str]): Request headers payload (Dict[str, str]): Request payload Returns: str: Markdown formatted image or error message """ try: # Submit the generation task response = requests.post( f"{self.valves.NOVITA_API_BASE_URL}/v3beta/flux-1-schnell", headers=headers, json=payload, ) response.raise_for_status() result = response.json() if not result.get("images"): return "Error: No images generated" # Get the first generated image image_data = result["images"][0] image_url = image_data["image_url"] # Convert image URL to base64 try: img_data = self.get_image_from_url(image_url) return f"" except Exception as e: print(f"Error fetching image: {e}") return f"Error fetching image directly. URL: {image_url}" except requests.exceptions.RequestException as e: return f"Error: Request failed: {e}" except Exception as e: return f"Error: {e}" def stream_response( self, headers: Dict[str, str], payload: Dict[str, Any] ) -> Generator[str, None, None]: """ Process a streaming image generation request. Args: headers (Dict[str, str]): Request headers payload (Dict[str, Any]): Request payload Yields: str: The generated image data """ yield self.non_stream_response(headers, payload) def pipes(self) -> List[Dict[str, str]]: """ Get the list of available pipes. Returns: List[Dict[str, str]]: The list of pipes """ return [{"id": "flux_schnell", "name": "FLUX.1 Schnell"}] def pipe( self, body: Dict[str, Any] ) -> Union[str, Generator[str, None, None], Iterator[str]]: """ Process the pipe request. Args: body (Dict[str, Any]): The request body Returns: Union[str, Generator[str, None, None], Iterator[str]]: The response """ headers = { "Authorization": f"Bearer {self.valves.NOVITA_API_KEY}", "Content-Type": "application/json", } prompt = get_last_user_message(body["messages"]) # Default payload configuration payload = { "prompt": prompt, "width": 1024, "height": 1024, "seed": random.randint(0, 4294967295), "steps": 4, "image_num": 1, "response_image_type": "png", } try: if body.get("stream", False): return self.stream_response(headers, payload) else: return self.non_stream_response(headers, payload) except Exception as e: return f"Error: {e}"