mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-11-15 12:13:17 +00:00
9b95e8b0b6
- do not provide automatically imported api functions and objects in `imaginairy` root module
- horrible hack to overcome horrible design choices by easy_install/setuptools
The hack modifies the installed script to remove the __import__ pkg_resources line
If we don't do this then the scripts will be slow to start up because of
pkg_resources.require() which is called by setuptools to ensure the
"correct" version of the package is installed.
before modification example:
```
__requires__ = 'imaginAIry==14.0.0b5'
__import__('pkg_resources').require('imaginAIry==14.0.0b5')
__file__ = '/home/user/projects/imaginairy/imaginairy/bin/aimg'
with open(__file__) as f:
exec(compile(f.read(), __file__, 'exec'))
```
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import torch
|
|
from torch.cuda import OutOfMemoryError
|
|
|
|
from imaginairy.api import imagine_image_files
|
|
from imaginairy.schema import ImaginePrompt
|
|
from imaginairy.utils import get_device
|
|
|
|
|
|
def assess_memory_usage():
|
|
assert get_device() == "cuda"
|
|
img_size = 3048
|
|
prompt = ImaginePrompt("strawberries", size=64, seed=1)
|
|
imagine_image_files([prompt], outdir="outputs")
|
|
datalog = []
|
|
while True:
|
|
torch.cuda.reset_peak_memory_stats()
|
|
prompt = ImaginePrompt(
|
|
"beautiful landscape, Unreal Engine 5, RTX, AAA Game, Detailed 3D Render, Cinema4D",
|
|
size=img_size,
|
|
seed=1,
|
|
steps=2,
|
|
)
|
|
try:
|
|
imagine_image_files([prompt], outdir="outputs")
|
|
except OutOfMemoryError as e:
|
|
print(f"Out of memory at {img_size}x{img_size} size image.")
|
|
print(e)
|
|
break
|
|
max_used = torch.cuda.max_memory_allocated() / 1024**3
|
|
datalog.append((img_size, max_used))
|
|
print(f"{img_size},{max_used:.2f}\n")
|
|
img_size += 128
|
|
|
|
with open("img_size_memory_usage.csv", "w", encoding="utf-8") as f:
|
|
f.write("img_size,max_used\n")
|
|
for img_size, max_used in datalog:
|
|
f.write(f"{img_size},{max_used:.2f}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
assess_memory_usage()
|