Made a few adjustments, please check out. 'samples' issue and batch splitting.

[sona9soba]

So... I have been having some issues and tried to find my own way of solving them.

1. samples argument in chroma.sample:

• I wanted to produce more than one protein outputs from chroma.sample(...) with 'samples' argument, but sadly, it only generated 'one', probably due to lack of actual code that takes 'samples' integer as the number of outputs.
• For my own use, I only needed to implement the 'design' part (since I wanted multiple sequence designs per one backbone sample). Thus, I simply added for loop in the chroma.sample() part, as seen in the models.chroma module script file.
• I made the minimal amount of changes since I did not want to ruin the logic of the model. So... I hope the team somehow fix this issue.

2. The sample run consumes TOO MUCH GPU memory.

• I use GeForce rtx3050ti(4GB memory), which is, honestly not the best GPU option out there. However, I had absolutely no problems running RFdiffusion, AlphaFold, and sometimes, even MD simulations with my tiny laptop. chroma module run, on the other hand, always fails with 'Segmentation Fault', due to lack of GPU memory available.
• In my case, it turned out to be the issue with the size of the batch. Step size as small as 150 causes segmentation fault issue. So, I made a small changes in the layers.structure.conditioner file and models.chroma file.
• For layers.structure, I added checkpoints for parts with intensive calculations. I added checkpoints to the conditioners I used(substructure, subsequence, procap) and all is well now.
• The more significant change though, is splitting the batch in models.chroma file. I simply split the number of steps to 100 steps, and avoided scoring part until the last batch, integrating the result of all batches for final outcomes. For smaller designs, batch size 50 works fine, but I used 100 for the sequence of 364 aa proteins.
• This is all I did, but it definitely helped. Again, I made minimal changes, since I did not want to ruin the system.

Please take this issues into your consideration when releasing the next version.

Thanks!
chroma_modifications.zip