test gradio

app.py

@@ -1,43 +1,45 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UNet2DConditionModel, AutoencoderKL, DDIMScheduler
+from diffusers import (
+    StableDiffusionControlNetPipeline,
+    ControlNetModel,
+    UNet2DConditionModel,
+    AutoencoderKL,
+    UniPCMultistepScheduler,
+)
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
-import os
 from huggingface_hub import login
+import os
 
-# Authenticate with Hugging Face
+# Log in to Hugging Face with token from environment variables
 token = os.getenv("HF_TOKEN")
 login(token=token)
 
-# 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 and other components
-controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float32)
-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,
+# Model and ControlNet IDs
+model_id = "runwayml/stable-diffusion-v1-5"  # Known compatible model with ControlNet
+controlnet_id = "lllyasviel/sd-controlnet-canny"  # ControlNet model for edge detection
+
+# Load ControlNet model and other components
+controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
+pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+    model_id,
     controlnet=controlnet,
-    scheduler=scheduler,
-    feature_extractor=feature_extractor
+    torch_dtype=torch.float16
 )
 
-pipeline = pipeline.to("cuda") if torch.cuda.is_available() else pipeline
+# Optional: Set up the faster scheduler
+pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
 
-# Define Gradio interface function
+# Enable CPU offloading for memory optimization
+pipeline.enable_model_cpu_offload()
+
+
+# Gradio interface function
 def generate_image(prompt, reference_image):
+    # Resize and prepare reference image
     reference_image = reference_image.convert("RGB").resize((512, 512))
+
+    # Generate image using the pipeline with ControlNet
     generated_image = pipeline(
         prompt=prompt,
         image=reference_image,
@@ -47,6 +49,7 @@ def generate_image(prompt, reference_image):
     ).images[0]
     return generated_image
 
+
 # Set up Gradio interface
 interface = gr.Interface(
     fn=generate_image,
@@ -55,8 +58,8 @@ interface = gr.Interface(
         gr.Image(type="pil", label="Reference Image (Style)")
     ],
     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 with ControlNet."
+    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)."
 )
 
 # Launch the Gradio interface