import torch
from diffusers import SD3Transformer2DModel, StableDiffusion3Pipeline
from fastapi import FastAPI, Body
# Efficient loading with 4-bit quantization
quantization_config = BitsAndBytesConfig(
   load_in_4bit=True, 
   bnb_4bit_quant_type="nf4", 
   bnb_4bit_compute_dtype=torch.bfloat16
)
# Pipeline setup
pipe = StableDiffusion3Pipeline.from_pretrained(
   repo_id,
   torch_dtype=torch.bfloat16,
   transformer=model_nf4
)
pipe.enable_xformers_memory_efficient_attention()
# FastAPI endpoint
@app.post("/generate")
def generate(
   prompt: str = Body(),
   negative_prompt: str = Body(),
   height: int = Body(),
   width: int = Body(), 
   num_inference_steps: int = Body(),
   guidance_scale: float = Body(),
   ):
   image_store = io.BytesIO()
   images = pipe(prompt=prompt, negative_prompt=negative_prompt, 
                height=height, width=width).images
   images[0].save(image_store, "PNG")
   return b64encode(image_store.getvalue())

Part 2: Langchain Integration

Next, we'll create a custom tool for Langchain:

from langchain.tools import BaseTool
class ImageGenerationTool(BaseTool):
   name = "image_generator"
   description = "Generate images from text descriptions"
   
   def _run(self, prompt: str) -> str:
       try:
           response = requests.post(
               "http://localhost:8000/generate",
               json={
                   "prompt": prompt,
                   "negative_prompt": "",
                   "height": 512,
                   "width": 512,
                   "num_inference_steps": 30,
                   "guidance_scale": 7.5
               }
           )
           return f"Image generated successfully: {response.text}"
       except Exception as e:
           return f"Error generating image: {str(e)}"

Part 3: Custom Langflow Component

Finally, let's create our custom Langflow component for visual workflow design:

class CustomComponent(Component):
   display_name = "Image generator"
   description = "Custom component to create image."
   icon = "custom_components"
   
   inputs = [
       MessageTextInput(
           name="input_prompt",
           required=True,
           display_name="Input Prompt"
       ),
       DropdownInput(
           name="input_height",
           options=["512", "1024"]
       ),
       Input(
           name="input_apikey", 
           password="true"
       )
   ]
   def build_output_message(self) -> Message:
       # API call setup
       response = requests.post(url, headers=headers, json=data)
       image_data = response.json()
       
       # File management
       timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
       filename = f"{timestamp}_{uuid.uuid4()}.jpg"
       
       # Save and return response
       return Message(
           text="image generated",
           content_blocks=[
               ContentBlock(
                   contents=[MediaContent(urls=[full_url_path])]
               )
           ]
       )