uuu

app.py

@@ -5,13 +5,15 @@ import gradio as gr
 import spaces
 from PIL import Image
 from huggingface_hub import login
-from diffusers.utils import load_image
 from torchvision import transforms
+from diffusers.utils import load_image
 
 from models.transformer_sd3 import SD3Transformer2DModel
 from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
 
-# Download IP Adapter if not exists
+# ----------------------------
+# Step 1: Download IP Adapter if not exists
+# ----------------------------
 url = "https://huggingface.co/InstantX/SD3.5-Large-IP-Adapter/resolve/main/ip-adapter.bin"
 file_path = "ip-adapter.bin"
 
@@ -24,22 +26,30 @@ if not os.path.exists(file_path):
                 file.write(chunk)
     print("Download completed!")
 
-# Hugging Face login
+# ----------------------------
+# Step 2: Hugging Face Login
+# ----------------------------
 token = os.getenv("HF_TOKEN")
+if not token:
+    raise ValueError("Hugging Face token not found. Set the 'HF_TOKEN' environment variable.")
 login(token=token)
 
-# Model paths
+# ----------------------------
+# Step 3: Model Paths
+# ----------------------------
 model_path = 'stabilityai/stable-diffusion-3.5-large'
 ip_adapter_path = './ip-adapter.bin'
 image_encoder_path = "google/siglip-so400m-patch14-384"
 
-# Load transformer and pipeline
+# ----------------------------
+# Step 4: Load Transformer and Pipeline
+# ----------------------------
 transformer = SD3Transformer2DModel.from_pretrained(
-    model_path, subfolder="transformer", torch_dtype=torch.bfloat16
+    model_path, subfolder="transformer", torch_dtype=torch.float16
 )
 
 pipe = StableDiffusion3Pipeline.from_pretrained(
-    model_path, transformer=transformer, torch_dtype=torch.bfloat16
+    model_path, transformer=transformer, torch_dtype=torch.float16
 ).to("cuda")
 
 pipe.init_ipadapter(
@@ -48,21 +58,29 @@ pipe.init_ipadapter(
     nb_token=64,
 )
 
-
-@spaces.GPU
-def gui_generation(prompt, ref_img, guidance_scale, ipadapter_scale):
-
-    # Load and preprocess the reference image
+# ----------------------------
+# Step 5: Image Preprocessing Function
+# ----------------------------
+def preprocess_image(image_path):
+    """Preprocess the input image for the pipeline."""
     preprocess = transforms.Compose([
         transforms.Resize((384, 384)),
         transforms.ToTensor(),
         transforms.ConvertImageDtype(torch.float16)
     ])
+    image = Image.open(image_path).convert('RGB')
+    return preprocess(image).unsqueeze(0).to("cuda")
 
-    ref_img = Image.open(ref_img.name).convert('RGB')
-    ref_img_tensor = preprocess(ref_img).unsqueeze(0).to("cuda")
+# ----------------------------
+# Step 6: Gradio Function
+# ----------------------------
+@spaces.GPU
+def gui_generation(prompt, ref_img, guidance_scale, ipadapter_scale):
+    """Generate an image using Stable Diffusion 3.5 Large with IP-Adapter."""
+    # Preprocess the reference image
+    ref_img_tensor = preprocess_image(ref_img.name)
 
-    # Generate the image
+    # Run the pipeline
     with torch.no_grad():
         image = pipe(
             width=1024,
@@ -78,8 +96,9 @@ def gui_generation(prompt, ref_img, guidance_scale, ipadapter_scale):
 
     return image
 
-
-# Set up Gradio interface
+# ----------------------------
+# Step 7: Gradio Interface
+# ----------------------------
 prompt_box = gr.Textbox(label="Prompt", placeholder="Enter your image generation prompt")
 ref_img = gr.File(label="Upload Reference Image")
 guidance_slider = gr.Slider(
@@ -108,4 +127,7 @@ interface = gr.Interface(
     description="Generates an image based on a text prompt and a reference image using Stable Diffusion 3.5 Large with IP-Adapter."
 )
 
-interface.launch(share=True)
+# ----------------------------
+# Step 8: Launch Gradio App
+# ----------------------------
+interface.launch(share=True)