test gradio

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import torch
 from diffusers import (
-    StableDiffusionControlNetPipeline,
+    StableDiffusion3Pipeline,  # For SD3 models like Stable Diffusion 3.5
     ControlNetModel,
-    UNet2DConditionModel,
+    SD3Transformer2DModel,  # Replacing UNet with SD3 transformer
     AutoencoderKL,
     UniPCMultistepScheduler,
 )
@@ -21,18 +21,18 @@ controlnet_id = "lllyasviel/control_v11p_sd15_inpaint"
 
 # Load each model component required by the pipeline
 controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
-unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", torch_dtype=torch.float16)
+transformer = SD3Transformer2DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
 vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
 feature_extractor = CLIPFeatureExtractor.from_pretrained(model_id)
 text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder")
 tokenizer = CLIPTokenizer.from_pretrained(model_id)
 
 # Initialize the pipeline with all components
-pipeline = StableDiffusionControlNetPipeline(
+pipeline = StableDiffusion3Pipeline(
+    transformer=transformer,  # Using SD3 transformer
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
-    unet=unet,
     controlnet=controlnet,
     scheduler=UniPCMultistepScheduler.from_config({"name": "UniPCMultistepScheduler"}),
     feature_extractor=feature_extractor,
@@ -69,7 +69,7 @@ interface = gr.Interface(
     ],
     outputs="image",
     title="Image Generation with ControlNet (Reference-Only Style Transfer)",
-    description="Generates an image based on a text prompt and style reference image using Stable Diffusion and ControlNet (reference-only mode)."
+    description="Generates an image based on a text prompt and style reference image using Stable Diffusion 3.5 and ControlNet (reference-only mode)."
 )
 
 # Launch the Gradio interface