test gradio

app.py

@@ -1,33 +1,43 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UNet2DConditionModel, AutoencoderKL, DDIMScheduler
+from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 import os
 from huggingface_hub import login
 
-# Log in with your Hugging Face token (assumed stored in HF_TOKEN)
+# Authenticate with Hugging Face
 token = os.getenv("HF_TOKEN")
 login(token=token)
 
-# Model IDs for the base Stable Diffusion model and ControlNet variant
-model_id = "stabilityai/stable-diffusion-3.5-large-turbo"
-controlnet_id = "lllyasviel/control_v11p_sd15_inpaint"  # Make sure this ControlNet is compatible
+# Define model and controlnet IDs
+model_id = "runwayml/stable-diffusion-v1-5"  # Use a fully compatible model
+controlnet_id = "lllyasviel/control_v11p_sd15_inpaint"  # ControlNet variant
 
-# Load ControlNet model and pipeline
+# Load ControlNet and other components
 controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float32)
-pipeline = StableDiffusionControlNetPipeline.from_pretrained(
-    model_id,
+unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
+vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae")
+scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
+text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32")
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
+
+# Initialize the pipeline with all required components
+pipeline = StableDiffusionControlNetPipeline(
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    unet=unet,
     controlnet=controlnet,
-    torch_dtype=torch.float32
+    scheduler=scheduler,
+    feature_extractor=feature_extractor
 )
-pipeline = pipeline.to("cuda") if torch.cuda.is_available() else pipeline
 
+pipeline = pipeline.to("cuda") if torch.cuda.is_available() else pipeline
 
-# Define the Gradio interface function
+# Define Gradio interface function
 def generate_image(prompt, reference_image):
-    # Ensure the reference image is in the correct format
     reference_image = reference_image.convert("RGB").resize((512, 512))
-
-    # Generate the image with ControlNet
     generated_image = pipeline(
         prompt=prompt,
         image=reference_image,
@@ -37,7 +47,6 @@ def generate_image(prompt, reference_image):
     ).images[0]
     return generated_image
 
-
 # Set up Gradio interface
 interface = gr.Interface(
     fn=generate_image,
@@ -47,7 +56,7 @@ interface = gr.Interface(
     ],
     outputs="image",
     title="Image Generation with Reference-Only Style Transfer",
-    description="Generate an image based on a text prompt and style reference image using Stable Diffusion 3.5 with ControlNet (reference-only mode)."
+    description="Generate an image based on a text prompt and style reference image using Stable Diffusion with ControlNet."
 )
 
 # Launch the Gradio interface