Update app.py

app.py

@@ -3,21 +3,23 @@ import torch
 from diffusers import AutoPipelineForText2Image
 
 # 2. Load the model
-# We use AutoPipelineForText2Image to automatically load the correct pipeline for SDXL.
-# torch_dtype=torch.float16 is a memory optimization.
-# variant="fp16" tells the model to use a smaller, faster version of the weights.
-# .to("cuda") moves the model to the GPU for fast inference.
+# We remove the .to("cuda") from the end of this line.
 print("Loading model...")
 pipe = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", 
     torch_dtype=torch.float16, 
     variant="fp16", 
     use_safetensors=True
-).to("cuda")
-print("Model loaded.")
+)
+print("Model partially loaded.")
 
-# 3. Define the image generation function
-# This function will take a text prompt and return a generated image.
+# ### CHANGE HERE ###
+# This is the magic line that enables memory-efficient offloading.
+pipe.enable_model_cpu_offload()
+print("Model loaded with CPU offloading.")
+
+
+# 3. Define the image generation function (This part remains the same)
 def generate_image(prompt, negative_prompt, steps, guidance):
     print(f"Generating image for prompt: {prompt}")
     # The 'pipe' object does all the work. We pass it the prompt and other parameters.
@@ -33,30 +35,23 @@ def generate_image(prompt, negative_prompt, steps, guidance):
     # The output is a PIL Image object, which Gradio can display directly.
     return image
 
-# 4. Create the Gradio interface
-# This is where we design the web UI.
+# 4. Create the Gradio interface (This part remains the same)
 with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky")) as demo:
     gr.Markdown("# 🖼️ Stable Diffusion XL Text-to-Image")
     gr.Markdown("Enter a text prompt and see the magic of AI-powered image generation!")
     
     with gr.Row():
         with gr.Column(scale=4):
-            # Textbox for the main prompt
             prompt_input = gr.Textbox(label="Your Prompt", placeholder="An astronaut riding a horse on Mars, cinematic, 4k")
-            # Textbox for the negative prompt
             negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, watermark, text")
-            # Submit button
             submit_btn = gr.Button("Generate Image", variant="primary")
         
         with gr.Column(scale=1):
-            # Sliders for advanced options
             steps_slider = gr.Slider(minimum=10, maximum=100, value=25, step=1, label="Inference Steps")
             guidance_slider = gr.Slider(minimum=0, maximum=20, value=7.5, step=0.1, label="Guidance Scale")
     
-    # Image component to display the output
     output_image = gr.Image(label="Generated Image", width=768, height=768)
 
-    # Define some example prompts to make it easy for users to start
     gr.Examples(
         examples=[
             ["A majestic lion wearing a crown, photorealistic", "cartoon, drawing", 30, 8],
@@ -66,9 +61,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky"))
         inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider]
     )
 
-    # 5. Connect the button to the function
-    # When the submit button is clicked, it will call the `generate_image` function
-    # with the values from the input components. The result will be displayed in `output_image`.
+    # 5. Connect the button to the function (This part remains the same)
     submit_btn.click(
         fn=generate_image,
         inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider],