Fix ZeroGPU/diffusers-drift breakage: controlnet channel mismatch, import paths, CUDA fork

Root cause: current diffusers dropped extra_condition_channels support in
FluxControlNetModel (breaks alimama-creative FLUX.1-dev-Controlnet-Inpainting-Beta,
which needs 68 input channels not 64), moved zero_module/BaseOutput import paths,
and changed FluxAttnProcessor2_0's rotary-embedding format to a (cos, sin) tuple.

Fixes:
- Wire app.py to the vendored controlnet_flux.FluxControlNetModel and
pipeline_flux_controlnet_inpaint.FluxControlNetInpaintingPipeline (already in
the repo, already used by main.py) instead of diffusers built-ins.
- Fix diffusers.models.controlnet -> diffusers.models.controlnets.controlnet import.
- Import spaces before torch (avoids CUDA init in parent process poisoning the
ZeroGPU worker fork).
- Drop unused optimum-quanto import/dependency (its CUDA extension import was
triggering a real cuInit in the parent process).
- Fix run_pipe() call signature mismatch in inpaint() (pre-existing bug).
- Swap vendored EmbedND rope for diffusers FluxPosEmbed (matching tuple format)
and strip the 3D batch dim from txt_ids/img_ids before pos_embed, mirroring
diffusers upstream, since FluxPosEmbed expects 2D ids.

app.py

@@ -1,14 +1,15 @@
+import spaces
 import gradio as gr
 import torch
-import spaces
-from diffusers import AutoencoderKL, FluxTransformer2DModel, FluxControlNetModel, FluxControlNetInpaintPipeline
+from diffusers import AutoencoderKL
 from diffusers.utils import load_image
 from transformer_flux import FluxTransformer2DModel
+from controlnet_flux import FluxControlNetModel
+from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
 from transformers import T5EncoderModel, CLIPTextModel
 from PIL import Image, ImageDraw
 import numpy as np
 from huggingface_hub import hf_hub_download
-from optimum.quanto import freeze, qfloat8, quantize
 
 
 controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
@@ -16,7 +17,7 @@ transformer = FluxTransformer2DModel.from_pretrained(
     "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dtype=torch.bfloat16
 )
 
-pipe = FluxControlNetInpaintPipeline.from_pretrained(
+pipe = FluxControlNetInpaintingPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     transformer=transformer,
     controlnet=controlnet,
@@ -165,19 +166,14 @@ def inpaint(image, width, height, overlap_percentage, num_inference_steps, resiz
 
     #generator = torch.Generator(device="cuda").manual_seed(42)
     result = run_pipe(
-        prompt=final_prompt,
+        final_prompt=final_prompt,
         height=height,
         width=width,
-        control_image=cnet_image,
-        control_mask=mask,
+        cnet_image=cnet_image,
+        mask=mask,
         num_inference_steps=num_inference_steps,
-        #generator=generator,
-        controlnet_conditioning_scale=0.9,
-        guidance_scale=3.5,
-        negative_prompt="",
-        true_guidance_scale=3.5,
-    ).images[0]
-   
+    )
+
     result = result.convert("RGBA")
     cnet_image.paste(result, (0, 0), mask)
 

controlnet_flux.py

@@ -15,14 +15,15 @@ from diffusers.utils import (
     scale_lora_layers,
     unscale_lora_layers,
 )
-from diffusers.models.controlnet import BaseOutput, zero_module
+from diffusers.utils import BaseOutput
+from diffusers.models.controlnets.controlnet import zero_module
 from diffusers.models.embeddings import (
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepTextProjEmbeddings,
 )
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_flux import FluxPosEmbed
 from transformer_flux import (
-    EmbedND,
     FluxSingleTransformerBlock,
     FluxTransformerBlock,
 )
@@ -59,9 +60,7 @@ class FluxControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         self.out_channels = in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
 
-        self.pos_embed = EmbedND(
-            dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
-        )
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
         text_time_guidance_cls = (
             CombinedTimestepGuidanceTextProjEmbeddings
             if guidance_embeds
@@ -293,8 +292,11 @@ class FluxControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         )
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
-        txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
-        ids = torch.cat((txt_ids, img_ids), dim=1)
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)
 
         block_samples = ()

requirements.txt

@@ -4,5 +4,4 @@ transformers
 safetensors
 accelerate
 sentencepiece
-peft
-optimum-quanto
+peft

transformer_flux.py

@@ -32,6 +32,7 @@ from diffusers.models.embeddings import (
     CombinedTimestepTextProjEmbeddings,
 )
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_flux import FluxPosEmbed
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -293,9 +294,7 @@ class FluxTransformer2DModel(
             self.config.num_attention_heads * self.config.attention_head_dim
         )
 
-        self.pos_embed = EmbedND(
-            dim=self.inner_dim, theta=10000, axes_dim=axes_dims_rope
-        )
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
         text_time_guidance_cls = (
             CombinedTimestepGuidanceTextProjEmbeddings
             if guidance_embeds
@@ -417,8 +416,11 @@ class FluxTransformer2DModel(
         )
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
-        txt_ids = txt_ids.expand(img_ids.size(0), -1, -1)
-        ids = torch.cat((txt_ids, img_ids), dim=1)
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)
 
         for index_block, block in enumerate(self.transformer_blocks):