DEV Community

kk1123
kk1123

Posted on

prepare python environment

Create virtual environment

  1. Create a project folder
  2. Cd to the project folder
  3. Create virtual environment: conda create -n ai_painting python=3.10
  4. Activate virtual environment: conda activate ai_painting

Install librarys

  1. Install CPU torch pip install torch trochvision trochaudio
  2. Install AI painting libraries pip install diffusers["torch"] transformers accelerate
  3. Install web framework:pip install "fastapi[standard]"

Test Installing

Create a python file:

import torch
from diffusers import StableDiffusionPipeline
import gc  # Garbage collector

print("๐Ÿ”ง Setting up optimized AI system...")

# Force CPU and clear memory
device = "cpu"
torch.cuda.empty_cache() if torch.cuda.is_available() else None
gc.collect()

print("๐Ÿš€ Loading model with memory optimization...")

try:
    # Load with memory optimizations
    model_id = "runwayml/stable-diffusion-v1-5"

    # Load to CPU with specific settings for stability
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float32,  # Use float32 for CPU stability
        use_safetensors=True
    )
    pipe = pipe.to(device)

    # Disable the safety checker to save memory
    pipe.safety_checker = None
    pipe.requires_safety_checker = False

    print("โœ… Model loaded! Generating smaller test image...")

    # Generate a smaller image to save memory
    prompt = "a simple red apple on a table"

    with torch.no_grad():
        image = pipe(
            prompt,
            num_inference_steps=20,  # Fewer steps = less memory
            guidance_scale=7.5,
            width=256,  # Smaller image
            height=256   # Smaller image
        ).images[0]

    image.save("test_small.jpg")
    print("๐Ÿ“ Success! Image saved as 'test_small.jpg'")

except Exception as e:
    print(f"โŒ Error: {e}")
Enter fullscreen mode Exit fullscreen mode

Top comments (0)