test gradio

app.py

@@ -2,27 +2,32 @@ import gradio as gr
 import torch
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 import os
-from huggingface_hub import HfApi, login
+from huggingface_hub import login
 
+# Log in with your Hugging Face token (assumed stored in HF_TOKEN)
 token = os.getenv("HF_TOKEN")
-login(token=token)  # Logs in with the token in Hugging Face Spaces
+login(token=token)
 
-# Load Stable Diffusion model and ControlNet reference-only model
+# 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"  # Use an appropriate ControlNet variant
+controlnet_id = "lllyasviel/control_v11p_sd15_inpaint"  # Make sure this ControlNet is compatible
 
+# Load ControlNet model and pipeline
 controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float32)
 pipeline = StableDiffusionControlNetPipeline.from_pretrained(
     model_id,
     controlnet=controlnet,
     torch_dtype=torch.float32
 )
+pipeline = pipeline.to("cuda") if torch.cuda.is_available() else pipeline
+
 
 # Define the Gradio interface function
 def generate_image(prompt, reference_image):
-    # Process reference image
-    reference_image = reference_image.resize((512, 512))
-    # Generate image with reference-only style transfer
+    # 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,
@@ -32,6 +37,7 @@ def generate_image(prompt, reference_image):
     ).images[0]
     return generated_image
 
+
 # Set up Gradio interface
 interface = gr.Interface(
     fn=generate_image,