Create optimized.py

optimized.py

@@ -0,0 +1,75 @@
+import torch
+import spaces
+import os
+from diffusers.utils import load_image
+from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
+import gradio as gr
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")
+
+# Load pipeline
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    torch_dtype=torch.bfloat16
+)
+pipe = FluxControlNetPipeline.from_pretrained(
+    "LPX55/FLUX.1-merged_uncensored",
+    controlnet=controlnet,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    vae=good_vae,
+    token=huggingface_token
+)
+# Add to your pipeline initialization:
+# pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_vae_slicing()  # Batch processing of VAE
+pipe.enable_model_cpu_offload()  # Use with accelerate
+
+# Convert all models to memory-efficient format
+pipe.to(memory_format=torch.channels_last)
+pipe.to("cuda")
+
+@spaces.GPU
+def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale):
+    # Load control image
+    control_image = control_image.resize((int(w * scale), int(h * scale)), PIL.Image.BICUBIC)
+    # control_image = load_image(control_image)
+    w, h = control_image.size
+    # Upscale x1
+    control_image = control_image.resize((int(w * scale), int(h * scale)))
+    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
+    image = pipe(
+        prompt=prompt,
+        control_image=control_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        height=control_image.size[1],
+        width=control_image.size[0],
+        torch_dtype=torch.bfloat16,
+        device_map="auto"
+    ).images[0]
+    torch.cuda.empty_cache()
+    return image
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.Slider(1, 3, value=1, label="Scale"),
+        gr.Slider(6, 30, value=8, label="Steps"),
+        gr.Image(type="pil", label="Control Image"),
+        gr.Slider(0, 1, value=0.6, label="ControlNet Scale"),
+        gr.Slider(1, 20, value=3.5, label="Guidance Scale"),
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Generated Image", format="png"),
+    ],
+    title="FLUX ControlNet Image Generation",
+    description="Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.",
+)
+
+# Launch the app
+iface.launch()